Update src/model.py

src/model.py

@@ -29,67 +29,15 @@ def Encoder(latent_channels=4):
         conv(64, latent_channels),
     )
 
-class DCAH(nn.Module):
-    def __init__(self, in_channels, embed_dim=64, dilation_rates=(1, 2, 4)):
-        super(DCAH, self).__init__()
-        self.in_channels = in_channels
-        self.embed_dim = embed_dim
-        self.dilated_convs = nn.ModuleList([
-            nn.Conv2d(in_channels, embed_dim, kernel_size=3, padding=rate, dilation=rate)
-            for rate in dilation_rates
-        ])
-        self.dilated_conv_merge = nn.Conv2d(embed_dim * len(dilation_rates), embed_dim, kernel_size=1)
-        self.query = nn.Conv2d(embed_dim, embed_dim, kernel_size=1)
-        self.key = nn.Conv2d(embed_dim, embed_dim, kernel_size=1)
-        self.value = nn.Conv2d(embed_dim, embed_dim, kernel_size=1)
-        self.refine = nn.Sequential(
-            nn.Conv2d(embed_dim, embed_dim, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.Conv2d(embed_dim, in_channels, kernel_size=1)
-        )
-
-    def forward(self, x):
-        dilated_features = [conv(x) for conv in self.dilated_convs]
-        concat_features = torch.cat(dilated_features, dim=1)
-        global_context = self.dilated_conv_merge(concat_features)
-        q = self.query(global_context)
-        k = self.key(global_context)
-        v = self.value(global_context)
-        attention = F.softmax(torch.matmul(q.flatten(2), k.flatten(2).transpose(-2, -1)), dim=-1)
-        attention_out = torch.matmul(attention, v.flatten(2)).view_as(global_context)
-        refined = self.refine(global_context + attention_out)
-        return refined
-
-def DecoderSeq(latent_channels=16): 
+def Decoder(latent_channels=4):
     return nn.Sequential(
-        Clamp(), 
-        conv(latent_channels, 48), 
-        nn.ReLU(),
-        Block(48, 48), Block(48, 48),  
-        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
-        Block(48, 48), Block(48, 48), 
-        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
-        Block(48, 48), 
-        nn.Upsample(scale_factor=2), conv(48, 48, bias=False),
-        Block(48, 48), 
-        conv(48, 3), 
+        Clamp(), conv(latent_channels, 64), nn.ReLU(),
+        Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),nn.ReLU(),
+        Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),nn.ReLU(),
+        Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),nn.ReLU(),
+        Block(64, 64), conv(64, 3),
     )
 
-
-class Decoder(nn.Module):
-    def __init__(self, latent_channels=16):
-        decoder = DecoderSeq(latent_channels=latent_channels)
-        refinement_head = DCAH(in_channels=3, embed_dim=64)
-        super(Decoder, self).__init__()
-        self.decoder = decoder
-        self.refinement_head = refinement_head
-
-    def forward(self, x):
-        decoded = self.decoder(x)
-        refined = self.refinement_head(decoded)
-        return refined
-        
-
 class Model(nn.Module):
     latent_magnitude = 3
     latent_shift = 0.5
@@ -103,16 +51,17 @@ class Model(nn.Module):
         if encoder_path is not None:
             encoder_state_dict = torch.load(encoder_path, map_location="cpu", weights_only=True)
             filtered_state_dict = {k.strip('encoder.'): v for k, v in encoder_state_dict.items() if k.strip('encoder.') in self.encoder.state_dict() and v.size() == self.encoder.state_dict()[k.strip('encoder.')].size()}
+            print(f" num of keys in filtered: {len(filtered_state_dict)} and in decoder: {len(self.encoder.state_dict())}")
             self.encoder.load_state_dict(filtered_state_dict, strict=False)
             
         if decoder_path is not None:
             decoder_state_dict = torch.load(decoder_path, map_location="cpu", weights_only=True)
-            filtered_state_dict = {k: v for k, v in decoder_state_dict.items() if k in self.decoder.state_dict() and v.size() == self.decoder.state_dict()[k].size()}
+            filtered_state_dict = {k.strip('decoder.'): v for k, v in decoder_state_dict.items() if k.strip('decoder.') in self.decoder.state_dict() and v.size() == self.decoder.state_dict()[k.strip('decoder.')].size()}
+            print(f" num of keys in filtered: {len(filtered_state_dict)} and in decoder: {len(self.decoder.state_dict())}")
             self.decoder.load_state_dict(filtered_state_dict, strict=False)
             
         self.encoder.requires_grad_(False)
-        self.decoder.decoder.requires_grad_(False)
-        self.decoder.refinement_head.requires_grad_(False)
+        self.decoder.requires_grad_(False)
 
     def guess_latent_channels(self, encoder_path):
         if "taef1" in encoder_path:return 16