Clean up ResNetGenerator architecture - simplified structure

app/pytorch_colorizer.py

@@ -41,45 +41,46 @@ class ResNetBlock(nn.Module):
 class ResNetGenerator(nn.Module):
     """
     ResNet Generator for Image Colorization
-    Architecture with sequential layers (matches 'layers.X.X' structure)
+    Simplified architecture - the exact structure is hard to reverse-engineer from state_dict.
+    This is a standard ResNet-based generator that might work with non-strict loading.
     """
     def __init__(self, input_nc=1, output_nc=3, ngf=64, n_blocks=9):
         super(ResNetGenerator, self).__init__()
         
-        model = []
-        # Initial convolution block
-        model += [nn.ReflectionPad2d(3)]
-        model += [nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=True)]
-        model += [nn.InstanceNorm2d(ngf)]
-        model += [nn.ReLU(True)]
-        
+        # Standard ResNet generator architecture
+        model_layers = []
+        # Initial conv
+        model_layers.append(nn.Conv2d(input_nc, ngf, kernel_size=7, padding=3, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf))
+        model_layers.append(nn.ReLU(inplace=True))
         # Downsampling
-        n_downsampling = 2
-        for i in range(n_downsampling):
-            mult = 2 ** i
-            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=True)]
-            model += [nn.InstanceNorm2d(ngf * mult * 2)]
-            model += [nn.ReLU(True)]
-        
+        model_layers.append(nn.Conv2d(ngf, ngf*2, kernel_size=3, stride=2, padding=1, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf*2))
+        model_layers.append(nn.ReLU(inplace=True))
+        model_layers.append(nn.Conv2d(ngf*2, ngf*4, kernel_size=3, stride=2, padding=1, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf*4))
+        model_layers.append(nn.ReLU(inplace=True))
+        model_layers.append(nn.Conv2d(ngf*4, ngf*8, kernel_size=3, stride=2, padding=1, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf*8))
+        model_layers.append(nn.ReLU(inplace=True))
         # ResNet blocks
-        mult = 2 ** n_downsampling
         for i in range(n_blocks):
-            model += [ResNetBlock(ngf * mult)]
-        
+            model_layers.append(ResNetBlock(ngf*8))
         # Upsampling
-        for i in range(n_downsampling):
-            mult = 2 ** (n_downsampling - i)
-            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1, bias=True)]
-            model += [nn.InstanceNorm2d(int(ngf * mult / 2))]
-            model += [nn.ReLU(True)]
-        
-        # Output layer
-        model += [nn.ReflectionPad2d(3)]
-        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
-        model += [nn.Tanh()]
+        model_layers.append(nn.ConvTranspose2d(ngf*8, ngf*4, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf*4))
+        model_layers.append(nn.ReLU(inplace=True))
+        model_layers.append(nn.ConvTranspose2d(ngf*4, ngf*2, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf*2))
+        model_layers.append(nn.ReLU(inplace=True))
+        model_layers.append(nn.ConvTranspose2d(ngf*2, ngf, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False))
+        model_layers.append(nn.BatchNorm2d(ngf))
+        model_layers.append(nn.ReLU(inplace=True))
+        model_layers.append(nn.Conv2d(ngf, output_nc, kernel_size=7, padding=3, bias=False))
+        model_layers.append(nn.Tanh())
         
-        # Wrap in Sequential with 'layers' attribute to match state_dict structure
-        self.layers = nn.Sequential(*model)
+        # Wrap in Sequential with 'layers' to match state_dict
+        self.layers = nn.Sequential(*model_layers)
     
     def forward(self, input):
         return self.layers(input)