trying skip connection

configs/direct_3.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '4'
+config['name'] = 'direct_3'
+config['notes'] = 'with curvature input and skip connection' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = True
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+config['input_features'] = 10
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256*2, 128, 9),
+    CNBlockConfig(128*2, 32, 9),
+    CNBlockConfig(32*2, 16, 3),
+    CNBlockConfig(16*2, 8, 3),
+    CNBlockConfig(8, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_curvature',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }

models/base.py

@@ -53,13 +53,13 @@ class AEEncoder(nn.Module):
         else:
             x = inputs
 
-        x = self.encoder(x)
+        x, features = self.encoder(x)
         x = self.act(x)
 
         if len(inputs) == 2:
-            return x, inputs[1]
+            return x, features, inputs[1]
         else:
-            return x
+            return x, features
 
 
 class Decoder(nn.Module):
@@ -75,12 +75,12 @@ class Decoder(nn.Module):
             self.softmax = nn.Softmax(dim=1)
 
     def forward(self, inputs):
-        if len(inputs) == 2:
+        if len(inputs) == 3:
-            x, mask = inputs
+            x, features, mask = inputs
         else:
-            x = inputs
+            x, features = inputs
 
-        x = self.decoder(x)
+        x = self.decoder([x, features])
 
         if self.is_label:
             x = x * mask[:,None,:,:]

models/convnext_orig.py

@@ -81,8 +81,6 @@ class ConvNeXtBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        layers_list: List[nn.Module] = []
-
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
@@ -95,7 +93,8 @@ class ConvNeXtBase(nn.Module):
                 activation_layer=None,
                 bias=True,
             ).to(device)
-        layers_list.append(self.stem)
+        self.stages = []
+        self.downsamples = []
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
@@ -107,17 +106,20 @@ class ConvNeXtBase(nn.Module):
                 sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                 stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
                 stage_block_id += 1
-            layers_list.append(nn.Sequential(*stage).to(device))
+            self.stages.append(nn.Sequential(*stage).to(device))
             if cnf.out_channels is not None:
                 # Downsampling
-                layers_list.append(
+                self.downsamples.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
                     ).to(device)
                 )
+        self.stages = nn.ModuleList(self.stages)
+        self.downsamples = nn.ModuleList(self.downsamples)
 
-        self.layers = nn.Sequential(*layers_list).to(device)
+        # self.layers = nn.Sequential(*layers_list).to(device)
+        # self.layers = nn.ModuleList(layers_list)
         self.act = nn.Tanh().to(device)
 
         # for m in self.modules():
@@ -127,9 +129,16 @@ class ConvNeXtBase(nn.Module):
         #             nn.init.zeros_(m.bias)
 
     def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.layers(x)
+        x = self.stem(x)
+        features = []
+        for i in range(len(self.stages)):
+            x = self.stages[i](x)
+            if i == len(self.stages) - 1:
+                continue
+            features.append(x)
+            x = self.downsamples[i](x)
         x = self.act(x)
-        return x
+        return x, features
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
@@ -165,13 +174,14 @@ class ConvNeXtReverseBase(nn.Module):
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
         self.stem = nn.Sequential(
-                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=4, stride=4),
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
                         norm_layer(block_setting[0].out_channels),
                     ).to(device)
         
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        layers_list: List[nn.Module] = []
+        self.stages = []
+        self.upsamples = []
 
 
         for i, cnf in enumerate(block_setting):
@@ -182,17 +192,17 @@ class ConvNeXtReverseBase(nn.Module):
                 sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                 stage.append(block(cnf.input_channels, layer_scale, sd_prob, kernel_size))
                 stage_block_id += 1
-            layers_list.append(nn.Sequential(*stage))
+            self.stages.append(nn.Sequential(*stage).to(device))
 
             if i == 0:
-                layers_list.append(self.stem)
+                self.upsamples.append(self.stem)
             else:
                 # Upsampling
-                layers_list.append(
+                self.upsamples.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    )
+                    ).to(device)
                 )
 
         self.last = nn.Sequential(
@@ -201,7 +211,8 @@ class ConvNeXtReverseBase(nn.Module):
                         Permute([0, 3, 1, 2]),
                     ).to(device)
         
-        self.layers = nn.Sequential(*layers_list).to(device)
+        self.stages = nn.ModuleList(self.stages)
+        self.upsamples = nn.ModuleList(self.upsamples)
         self.activation = nn.Tanh().to(device) if not is_label else nn.Identity().to(device)
 
 
@@ -213,7 +224,15 @@ class ConvNeXtReverseBase(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        x = self.layers(x)
+        x, features = x
+        for i in range(len(self.stages)-1):
+            if i == 0:
+                x = self.stages[i](x)
+            else:
+                x = self.stages[i](torch.cat([x, features[-i]], dim=1))
+            x = self.upsamples[i](x)
+        x = self.stages[-1](x)
+        x = self.upsamples[-1](x)
         x = self.last(x)
         x = self.activation(x)
         return x

scripts/direct_train.py

@@ -13,7 +13,7 @@ args = vars(parser.parse_args())
 config_module = importlib.import_module(args['config'])
 config = config_module.config
 if args['debug']:
-    from configs.direct_1 import config
+    from configs.direct_3 import config
     # from configs.ae_static_orig_2 import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']