From 5a1276f045df6d1dcad548d689d3485b499d6cb2 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 24 May 2023 17:46:13 +0200 Subject: [PATCH] trying skip connection --- configs/direct_3.py | 78 +++++++++++++++++++++++++++++++++++++++++ models/base.py | 14 ++++---- models/convnext_orig.py | 51 ++++++++++++++++++--------- scripts/direct_train.py | 2 +- 4 files changed, 121 insertions(+), 24 deletions(-) create mode 100644 configs/direct_3.py diff --git a/configs/direct_3.py b/configs/direct_3.py new file mode 100644 index 0000000..149cd53 --- /dev/null +++ b/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', + } \ No newline at end of file diff --git a/models/base.py b/models/base.py index d058f0d..3e7c251 100644 --- a/models/base.py +++ b/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: - x, mask = inputs + if len(inputs) == 3: + 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,:,:] diff --git a/models/convnext_orig.py b/models/convnext_orig.py index b466bf5..105dfb2 100644 --- a/models/convnext_orig.py +++ b/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 diff --git a/scripts/direct_train.py b/scripts/direct_train.py index fe8c2b1..2dca17f 100644 --- a/scripts/direct_train.py +++ b/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']