diff --git a/tld/denoiser.py b/tld/denoiser.py index ac29dae..11310e3 100644 --- a/tld/denoiser.py +++ b/tld/denoiser.py @@ -19,14 +19,22 @@ def __init__(self, patch_size, img_size, embed_dim, dropout, n_layers, mlp_multi self.n_layers = n_layers self.mlp_multiplier = mlp_multiplier - seq_len = int((self.img_size/self.patch_size)*((self.img_size/self.patch_size))) + seq_len = img_size**2#int((self.img_size/self.patch_size)*((self.img_size/self.patch_size))) patch_dim = self.n_channels*self.patch_size*self.patch_size + self.first_proj = nn.Sequential(Rearrange('bs d h w -> bs (h w) d'), + nn.Linear(n_channels, self.embed_dim//4), + nn.LayerNorm(self.embed_dim//4), + ) + + self.cond_proj = nn.Linear(self.embed_dim, self.embed_dim//4) + self.patchify_and_embed = nn.Sequential( - nn.Conv2d(self.n_channels, patch_dim, kernel_size=self.patch_size, stride=self.patch_size), + Rearrange('bs (h w) d ->bs d h w', h = self.img_size), + nn.Conv2d(self.embed_dim//4, self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size), Rearrange('bs d h w -> bs (h w) d'), - nn.LayerNorm(patch_dim), - nn.Linear(patch_dim, self.embed_dim), + nn.LayerNorm(self.embed_dim), + nn.Linear(self.embed_dim, self.embed_dim), nn.LayerNorm(self.embed_dim) ) @@ -35,9 +43,15 @@ def __init__(self, patch_size, img_size, embed_dim, dropout, n_layers, mlp_multi p1=self.patch_size, p2=self.patch_size) - self.pos_embed = nn.Embedding(seq_len, self.embed_dim) + self.pos_embed = nn.Embedding(seq_len, self.embed_dim//4) self.register_buffer('precomputed_pos_enc', torch.arange(0, seq_len).long()) + self.first_decoder_block = DecoderBlock(embed_dim=self.embed_dim//4, ##aplied on priginal image + mlp_multiplier=self.mlp_multiplier, + is_causal=False, + dropout_level=self.dropout, + mlp_class=MLPSepConv) + self.decoder_blocks = nn.ModuleList([DecoderBlock(embed_dim=self.embed_dim, mlp_multiplier=self.mlp_multiplier, #note that this is a non-causal block since we are @@ -52,13 +66,24 @@ def __init__(self, patch_size, img_size, embed_dim, dropout, n_layers, mlp_multi def forward(self, x, cond): - x = self.patchify_and_embed(x) + x = self.first_proj(x) ### 4,h,w -> h*w, 4 -> h*w, d/4 + pos_enc = self.precomputed_pos_enc[:x.size(1)].expand(x.size(0), -1) x = x+self.pos_embed(pos_enc) + ##have to project cond: + cond_projected = self.cond_proj(cond) + x = self.first_decoder_block(x, cond_projected) ###same dim as original data in and out. + x = self.patchify_and_embed(x) ## h*w, d/4 -> d/4, h, w -> h/2*w/2, d + for block in self.decoder_blocks: x = block(x, cond) + ## should I add another attention layer here at 1024 tokens? + + #h/2*w/2, d -> d, h/2, w/2 -> upsample -> attenion block -> h*w, d -> h*w, 4 -> + + return self.out_proj(x) class Denoiser(nn.Module): diff --git a/tld/train.py b/tld/train.py index 5d665ab..b5cee21 100644 --- a/tld/train.py +++ b/tld/train.py @@ -107,9 +107,13 @@ def main(config: ModelConfig, dataconfig: DataConfig): patch_size=config.patch_size, embed_dim=config.embed_dim, dropout=config.dropout, n_layers=config.n_layers) + loss_fn = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=config.lr) + accelerator.print("Compiling model:") + model = torch.compile(model) + if not config.from_scratch: accelerator.print("Loading Model:") wandb.restore(config.model_name, run_path=f"apapiu/cifar_diffusion/runs/{config.run_id}", diff --git a/tld/transformer_blocks.py b/tld/transformer_blocks.py index 076ce74..2f87ddb 100644 --- a/tld/transformer_blocks.py +++ b/tld/transformer_blocks.py @@ -30,15 +30,14 @@ def forward(self, q, k, v, attn_mask=None): assert q.size(-1) == k.size(-1) assert k.size(-2) == v.size(-2) - q, k, v = [rearrange(x, 'bs n (d h) -> bs h n d', h=self.n_heads) for x in [q,k,v]] - q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + q, k, v = [rearrange(x, 'bs n (h d) -> bs h n d', h=self.n_heads) for x in [q,k,v]] out = nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=self.is_causal, dropout_p=self.dropout_level if self.training else 0) - - out = rearrange(out, 'bs h n d -> bs n (d h)', h=self.n_heads) + + out = rearrange(out, 'bs h n d -> bs n (h d)', h=self.n_heads) return out @@ -102,14 +101,14 @@ class DecoderBlock(nn.Module): def __init__(self, embed_dim, is_causal, mlp_multiplier, dropout_level, mlp_class=MLP): super().__init__() self.self_attention = SelfAttention(embed_dim, is_causal, dropout_level, n_heads=embed_dim//64) - self.cross_attention = CrossAttention(embed_dim, is_causal=False, dropout_level=0, n_heads=4) + self.cross_attention = CrossAttention(embed_dim, is_causal=False, dropout_level=0, n_heads=embed_dim//64) self.mlp = mlp_class(embed_dim, mlp_multiplier, dropout_level) self.norm1 = nn.LayerNorm(embed_dim) self.norm2 = nn.LayerNorm(embed_dim) self.norm3 = nn.LayerNorm(embed_dim) def forward(self, x, y): - x = self.norm1(self.self_attention(x) + x) - x = self.norm2(self.cross_attention(x, y) + x) - x = self.norm3(self.mlp(x) + x) + x = self.self_attention(self.norm1(x)) + x + x = self.cross_attention(self.norm2(x), y) + x + x = self.mlp(self.norm3(x)) + x return x \ No newline at end of file