working static ae training

configs/ae_static.py

@@ -2,11 +2,13 @@ from utils.convnext_utils import CNBlockConfig
 
 config = dict()
 
-config['gpu'] = '5'
+config['gpu'] = '1'
 config['name'] = 'static_torch_1'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_1'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'
@@ -60,7 +62,6 @@ config['cond_ae']['activation'] = 'tanh'
 
 config['lr'] = 0.0005
 config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
 
 config['signature'] = { 'id',
                         'frame_nr',

configs/ae_static_orig.py

@@ -0,0 +1,71 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '2'
+config['name'] = 'static_torch_orig_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+
+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'] = True
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 27),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 27),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

configs/ae_static_orig_2.py

@@ -0,0 +1,71 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '3'
+config['name'] = 'static_torch_orig_2'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'static'] 
+
+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'] = True
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = False
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(8, 16, 3),
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 8, 3),
+    CNBlockConfig(8, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.0005
+config['cond_ae']['upscale_output'] = False
+config['pretrained_model'] = ''
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_static',
+                        'mesh_static',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

custom_utils/custom_utils.py

@@ -291,22 +291,6 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''):
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    if not disable_assert:
-        assert manager.latest_checkpoint is not None, checkpoint
-        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint
-    if assert_consumed_on:
-        ckpt.restore(manager.latest_checkpoint).assert_consumed()
-    else:
-        if iteration != '':
-            name_base = manager.latest_checkpoint.split('-')[0]
-            ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0]
-            ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed()
-        else:
-            ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()
-
 
 def get_summary_str(model):
     # from torchview import draw_graph

data/dataset.py

@@ -53,15 +53,15 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
 
 
 class StaticCache():
-    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False) -> None:
+    def __init__(self, folder, signature, resolutions=[512], is_cape=False, device='cuda') -> None:
         self.store = dict()
         self.signature = signature
         self.folder = folder
-        self.in_memory = in_memory
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
         self.resolutions = resolutions
         self.filled = ''
         self.is_cape = is_cape
+        self.device = device
         
         if is_cape:
             self.uv_masks = dict()
@@ -105,22 +105,14 @@ class StaticCache():
     def __getitem__(self, id) -> dict:
         if id not in self.store:
             item = self.read(id)
-            if not self.in_memory:
-                return item
-            else:
-                self.store[id] = item
+            stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
+            for key, value in stored_data.items():
+                if key.startswith('uv_') and value.shape[-1] == 3:
+                    stored_data[key] = torch.permute(value, [2,0,1])
+            self.store[id] = stored_data
         
-        return self.store[id]
-
-
-    def expand_batch(self, batch):
-        batch_size = batch['id'].shape[0]
-        items = []
-        for i in range(batch_size):
-            items.append(self.__getitem__(batch['id'][i].numpy().decode('UTF-8')))
-        for key in items[0].keys():
-            batch[key] = torch.stack([item[key] for item in items], dim=0)
-        return batch
+        data = { key: value.to(self.device) for key, value in self.store[id].items()}
+        return data
 
 
 class Dataset(Dataset):
@@ -148,7 +140,9 @@ class Dataset(Dataset):
         frame : Frame = self.frames[idx]
         if self.cache is not None:
             if frame in self.cache.keys():
-                return self.cache[frame]
+                data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
+                data = { **item, **self.static_cache[frame.id] }
+                return data
 
         separator = '_'
         if self.is_cape:
@@ -208,9 +202,8 @@ class Dataset(Dataset):
                 uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
                 item['uv_posed' + res_string] = uv_posed
 
-        data = { **item, **self.static_cache[frame.id] }
 
-        data = { key: (torch.from_numpy(value).to(self.device) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in data.items()}
+        data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
 
         for key, value in data.items():
             if key.startswith('uv_') and value.shape[-1] == 3:
@@ -218,6 +211,9 @@ class Dataset(Dataset):
 
         if self.cache is not None:
             self.cache[frame] = data
+        
+        data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
+        data = { **item, **self.static_cache[frame.id] }
         return data
 
 

models/base.py

@@ -69,7 +69,8 @@ class Decoder(nn.Module):
 
         self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True)
             
-        self.softmax = nn.Softmax(dim=1)
+        if self.is_label:
+            self.softmax = nn.Softmax(dim=1)
 
     def forward(self, inputs):
         if len(inputs) == 2:

models/convnext.py

@@ -8,25 +8,10 @@ from torch import nn, Tensor
 from torch.nn import functional as F
 
 from torchvision.ops.misc import Conv2dNormActivation, Permute
-# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
-# from torchvision.ops import Permute
 from torchvision.ops.stochastic_depth import StochasticDepth
 from utils.convnext_utils import CNBlockConfig
 
 
-__all__ = [
-    "ConvNeXt",
-    "ConvNeXt_Tiny_Weights",
-    "ConvNeXt_Small_Weights",
-    "ConvNeXt_Base_Weights",
-    "ConvNeXt_Large_Weights",
-    "convnext_tiny",
-    "convnext_small",
-    "convnext_base",
-    "convnext_large",
-]
-
-
 class LayerNorm2d(nn.LayerNorm):
     def forward(self, x: Tensor) -> Tensor:
         x = x.permute(0, 2, 3, 1)
@@ -91,11 +76,7 @@ class ConvNeXtBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        # self.permute_cf = Permute([0, 3, 1, 2])
-        # self.permute_cl = Permute([0, 2, 3, 1])
-
-        self.stages: List[nn.Module] = []
-        self.downsamples: List[nn.Module] = []
+        layers_list: List[nn.Module] = []
 
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
@@ -121,17 +102,17 @@ 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
-            self.stages.append(nn.Sequential(*stage).to(device))
+            layers_list.append(nn.Sequential(*stage))
             if cnf.out_channels is not None:
                 # Downsampling
-                self.downsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    ).to(device)
+                    )
                 )
 
-        self.act = nn.Tanh()
+        self.layers = nn.Sequential(*layers_list).to(device)
 
         for m in self.modules():
             if isinstance(m, (nn.Conv2d, nn.Linear)):
@@ -140,25 +121,15 @@ class ConvNeXtBase(nn.Module):
                     nn.init.zeros_(m.bias)
 
     def _forward_impl(self, x: Tensor) -> Tensor:
-        # x = self.permute_cf(x)
         x = self.stem(x)
-        # features: List[Tensor] = []
-        for i in range(len(self.stages) - 1):
-            x = self.stages[i](x)
-            # features.append(x)
-            x = self.downsamples[i](x)
-        x = self.stages[-1](x)
-        x = self.act(x)
-
-        # x = self.permute_cl(x)
+        x = self.layers(x)
         return x
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
-
-class ConvNeXtReverseBase2(nn.Module):
+class ConvNeXtReverseBase(nn.Module):
     def __init__(
         self,
         block_setting: List[CNBlockConfig],
@@ -187,21 +158,6 @@ class ConvNeXtReverseBase2(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        # self.permute_cf = Permute([0, 3, 1, 2])
-        # self.permute_cl = Permute([0, 2, 3, 1])
-
-        # Stem
-        # firstconv_output_channels = block_setting[0].input_channels
-        # self.first = nn.Sequential(
-        #             nn.ConvTranspose2d(
-        #                 firstconv_output_channels,
-        #                 firstconv_output_channels,
-        #                 kernel_size=2,
-        #                 stride=2,
-        #                 padding=0,
-        #                 ),
-        #             norm_layer(firstconv_output_channels),
-        #         ).to(device)
         self.first = nn.Sequential(
                         nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
                         norm_layer(block_setting[0].out_channels),
@@ -210,8 +166,8 @@ class ConvNeXtReverseBase2(nn.Module):
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        self.stages: List[nn.Module] = []
-        self.upsamples: List[nn.Module] = []
+        layers_list: List[nn.Module] = []
+        
         for cnf in block_setting[1:]:
             if cnf.input_channels is not None:
                 # Bottlenecks
@@ -221,32 +177,29 @@ class ConvNeXtReverseBase2(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
-                self.stages.append(nn.Sequential(*stage).to(device))
+                layers_list.append(nn.Sequential(*stage))
             else:
-                self.stages.append(nn.Identity().to(device))
+                layers_list.append(nn.Identity())
 
             if cnf.out_channels is not None:
                 # Upsampling
-                self.upsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    ).to(device)
+                    )
                 )
 
+        self.layers = nn.Sequential(*layers_list).to(device)
+
         self.last = nn.Conv2d(
                     block_setting[-1].input_channels,
                     3 if not is_label else nr_of_labels,
                     kernel_size=3,
                     padding='same',
                     ).to(device)
-        self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
         
-        # self.last = nn.Sequential(
-        #                 Permute([0, 2, 3, 1]),
-        #                 nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
-        #                 Permute([0, 3, 1, 2]),
-        #             ).to(device)
+        self.activation = nn.Tanh() if not is_label else nn.Softmax(dim=1)
 
 
         for m in self.modules():
@@ -257,134 +210,16 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        # input, features = x
-        # x = self.permute_cf(x)
-
         x = self.first(x)
-        for i in range(len(self.stages)):
-            # if self.use_skip:
-            #     x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
-            # else:
-            x = self.stages[i](x)
-            x = self.upsamples[i](x)
-        
-        # x = self.stages[-2](x)
-        # x = self.upsamples[-1](x)
-
-        # x = self.stages[-1](x)
+        x = self.layers(x)
         x = self.last(x)
         x = self.activation(x)
-        # x = self.permute_cl(x)
-        
         return x
 
     def forward(self, x) -> Tensor:
         return self._forward_impl(x)
 
 
-class ConvNeXtReverseBase(nn.Module):
-    def __init__(
-        self,
-        block_setting: List[CNBlockConfig],
-        stochastic_depth_prob: float = 0.0,
-        layer_scale: float = 1e-6,
-        block: Optional[Callable[..., nn.Module]] = None,
-        norm_layer: Optional[Callable[..., nn.Module]] = None
-    ) -> None:
-        super().__init__()
-
-        if not block_setting:
-            raise ValueError("The block_setting should not be empty")
-        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
-            raise TypeError("The block_setting should be List[CNBlockConfig]")
-
-        if block is None:
-            block = CNBlock
-
-        if norm_layer is None:
-            norm_layer = partial(LayerNorm2d, eps=1e-6)
-
-        layers: List[nn.Module] = []
-
-
-        # add bottleneck which will be executed at the end, to process after the final larger upsampling
-
-
-
-        # Stem
-        firstconv_output_channels = block_setting[0].input_channels
-
-        layers.append(
-            nn.Sequential(
-                nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
-
-            )
-        )
-
-        stage: List[nn.Module] = []
-        for _ in range(3):
-            stage.append(block(16, layer_scale, 0.0))
-        layers.append(nn.Sequential(*stage))
-        
-        layers.append(
-            # Convtr NormActivation(
-            #     firstconv_output_channels,
-            #     3,
-            #     kernel_size=4,
-            #     stride=4,
-            #     padding=0,
-            #     norm_layer=None,
-            #     activation_layer=None,
-            #     bias=True,
-            # )
-            nn.Sequential(
-                nn.ConvTranspose2d(
-                    firstconv_output_channels,
-                    16,
-                    kernel_size=4,
-                    stride=4,
-                    padding=0,
-                    ),
-            )
-        )
-
-        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
-        stage_block_id = 0
-        for cnf in block_setting:
-            # Bottlenecks
-            stage: List[nn.Module] = []
-            for _ in range(cnf.num_layers):
-                # adjust stochastic depth probability based on the depth of the stage block
-                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
-                stage.append(block(cnf.input_channels, layer_scale, sd_prob))
-                stage_block_id += 1
-            layers.append(nn.Sequential(*stage))
-            if cnf.out_channels is not None:
-                # Upsampling
-                layers.append(
-                    nn.Sequential(
-                        norm_layer(cnf.out_channels),
-                        nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
-                    )
-                )
-
-        self.features = nn.Sequential(*reversed(layers))
-
-        for m in self.modules():
-            if isinstance(m, (nn.Conv2d, nn.Linear)):
-                nn.init.trunc_normal_(m.weight, std=0.02)
-                if m.bias is not None:
-                    nn.init.zeros_(m.bias)
-
-    def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.features(x)
-        # x = torch.tanh(x)
-        return x
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self._forward_impl(x)
-
-
 def _convnext(
     block_setting: List[CNBlockConfig],
     stochastic_depth_prob: float,
@@ -397,7 +232,7 @@ def _convnext(
 ) -> ConvNeXtBase:
 
     if reversed:
-        model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
     else:
         model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)
 

models/convnext_orig.py

@@ -0,0 +1,259 @@
+# taken from torchvision.models.convnext
+
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
+# from torchvision.ops import Permute
+from torchvision.ops.stochastic_depth import StochasticDepth
+from utils.convnext_utils import CNBlockConfig
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+class CNBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        layer_scale: float = 1e-6,
+        stochastic_depth_prob: float = 0.0,
+        kernel_size: int = 7,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.block = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=kernel_size, padding='same', groups=dim, bias=True),
+            Permute([0, 2, 3, 1]),
+            norm_layer(dim),
+            nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+            nn.GELU(),
+            nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+            Permute([0, 3, 1, 2]),
+        )
+        self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.layer_scale * self.block(input)
+        result = self.stochastic_depth(result)
+        result += input
+        return result
+
+
+class ConvNeXtBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        kernel_size: int = 7,
+    ) -> None:
+        super().__init__()
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        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(
+                3,
+                firstconv_output_channels,
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                norm_layer=norm_layer,
+                activation_layer=None,
+                bias=True,
+            ).to(device)
+        layers_list.append(self.stem)
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        for cnf in block_setting:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                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))
+            if cnf.out_channels is not None:
+                # Downsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    ).to(device)
+                )
+
+        self.layers = nn.Sequential(*layers_list).to(device)
+        self.act = nn.Tanh().to(device)
+
+        # for m in self.modules():
+        #     if isinstance(m, (nn.Conv2d, nn.Linear)):
+        #         nn.init.trunc_normal_(m.weight, std=0.02)
+        #         if m.bias is not None:
+        #             nn.init.zeros_(m.bias)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.layers(x)
+        x = self.act(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+class ConvNeXtReverseBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip: bool = False,
+        is_label: bool = False,
+        nr_of_labels: Optional[int] = None,
+        kernel_size: int = 7,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            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),
+                        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] = []
+
+
+        for i, cnf in enumerate(block_setting):
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                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))
+
+            if i == 0:
+                layers_list.append(self.stem)
+            else:
+                # Upsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.last = nn.Sequential(
+                        Permute([0, 2, 3, 1]),
+                        nn.Linear(in_features=block_setting[-1].out_channels, out_features=3 if not is_label else nr_of_labels, bias=True),
+                        Permute([0, 3, 1, 2]),
+                    ).to(device)
+        
+        self.layers = nn.Sequential(*layers_list).to(device)
+        self.activation = nn.Tanh().to(device) if not is_label else nn.Identity().to(device)
+
+
+        # for m in self.modules():
+        #     if isinstance(m, (nn.Conv2d, nn.Linear)):
+        #         nn.init.trunc_normal_(m.weight, std=0.02)
+        #         if m.bias is not None:
+        #             nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        x = self.layers(x)
+        x = self.last(x)
+        x = self.activation(x)
+        return x
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+
+def _convnext(
+    block_setting: List[CNBlockConfig],
+    stochastic_depth_prob: float,
+    reversed: bool,
+    is_label: bool,
+    nr_of_labels: Optional[int],
+    kernel_size: int,
+    device='cuda',
+    **kwargs: Any,
+) -> ConvNeXtBase:
+
+    if reversed:
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
+    else:
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)
+
+
+    return model
+
+
+class ConvNeXt(nn.Module):
+    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, is_label=False, nr_of_labels=None, kernel_size=7, device='cuda', **kwargs) -> None:
+        super().__init__()
+
+        self.model = _convnext(block_setting, stochastic_depth_prob, reversed, is_label, nr_of_labels, kernel_size, device=device, **kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.model(x)
+        return x
+        
+
+
+if __name__ == '__main__':
+    model = ConvNeXt().to('cuda')
+    test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
+
+    pass

scripts/ae_static_train.py

@@ -13,9 +13,8 @@ import wandb
 import random
 import numpy as np
 import torch
-from torch import nn
-from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
-from custom_utils.custom_utils import UV2Mesh, get_summary_html, load_model, maskFacesBatch, get_summary_str
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str
 import plotly.graph_objects as go
 import plotly.express as ex
 import torch.optim as optim
@@ -26,15 +25,18 @@ else:
     from losses.losses import edge_loss
 
 from data.dataset import Dataset
-from torch.utils.data import DataLoader
-from custom_utils.custom_utils import batch_to_dict, plot_3d
+# from torch.utils.data import DataLoader
+from custom_utils.custom_utils import plot_3d
 from models.cond_encoder import AutoEncoder
-from utils.checkpoints import CheckpointSaver
+from utils.checkpoints import CheckpointSaver, load_model
 from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+
+from pytorch_msssim import MS_SSIM
 
 
 def init_wandb():
-    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp/wandb', tags=config['tags'])
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/ae_static_train.py')
     artifact.add_file('losses/losses.py')
@@ -57,19 +59,21 @@ def train_step(batch):
     model.train()
     pred = model([batch['uv_static'], batch['uv_mask']])
 
-    mesh_pred = uv2mesh(pred)
-    mesh_orig = batch['mesh_static']
-
-    e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
-    l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
     ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
     loss = 1e3 * (1 - ms_ssim)
+    # loss = l1_loss
 
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()
 
-    metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    model.eval()
+    with torch.inference_mode():
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_static']
+        l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1))
+        e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
     
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
         
@@ -151,8 +155,6 @@ def train():
             ckpt_saver_enc(model.encoder, step, optimizer)
             ckpt_saver_dec(model.decoder, step, optimizer)
 
-        print(step)
-
         step += 1
 
 
@@ -170,24 +172,25 @@ if __name__ == '__main__':
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
-    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'])
-    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'])
-    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=False)
-    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=False)
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
     model : AutoEncoder = AutoEncoder(config['cond_ae'])
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
 
-    _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
 
     if config['pretrained_model'] != '':
         load_model(model.encoder, config['pretrained_model'] + '_enc')
         load_model(model.decoder, config['pretrained_model'] + '_dec')
 
-    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=20)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=20)
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=5)
 
     run_name = config['name']
     print(f'started training: {run_name}')

utils/checkpoints.py

@@ -5,7 +5,7 @@ import logging
 
 
 class CheckpointSaver:
-    def __init__(self, dirpath, decreasing=False, top_n=5):
+    def __init__(self, dirpath, top_n=5):
         """
         dirpath: Directory path where to store all model weights 
         decreasing: If decreasing is `True`, then lower metric is better
@@ -14,25 +14,27 @@ class CheckpointSaver:
         if not os.path.exists(dirpath): os.makedirs(dirpath)
         self.dirpath = dirpath
         self.top_n = top_n
-        self.decreasing = decreasing
         self.top_model_paths = []
-        self.best_metric_val = np.Inf if decreasing else -np.Inf
         
-    def __call__(self, model, epoch, optimizer, metric_val=np.Inf):
-        model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_epoch{epoch}.pt')
-        save = metric_val<self.best_metric_val if self.decreasing else metric_val>self.best_metric_val
-        if save: 
-            logging.info(f"Current metric value better than {metric_val} better than best {self.best_metric_val}, saving model at {model_path}")
-            self.best_metric_val = metric_val
-            torch.save({ 'model': model.state_dict(), 'optimizer':optimizer }, model_path)
-            self.top_model_paths.append({'path': model_path, 'score': metric_val})
-            self.top_model_paths = sorted(self.top_model_paths, key=lambda o: o['score'], reverse=not self.decreasing)
+    def __call__(self, model, iteration, optimizer):
+        model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt')
+        torch.save({ 'model': model.state_dict(), 'optimizer': optimizer }, model_path)
+        self.top_model_paths.append({'path': model_path})
         if len(self.top_model_paths)>self.top_n: 
             self.cleanup()
     
     def cleanup(self):
-        to_remove = self.top_model_paths[self.top_n:]
+        to_remove = self.top_model_paths[:len(self.top_model_paths)-self.top_n]
         logging.info(f"Removing extra models.. {to_remove}")
         for o in to_remove:
             os.remove(o['path'])
-        self.top_model_paths = self.top_model_paths[:self.top_n]
+        self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]
+
+
+def load_model(model, checkpoint, optimizer=None, iteration=''):
+    # list of checkpoint files in the folder checkpoints
+    checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
+    checkpoint = torch.load(checkpoints[-1])
+    model.load_state_dict(checkpoint['model'])
+    if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'])
+

utils/utils.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import numpy as np
+from torch.utils.data import DataLoader
 
 
 
@@ -25,3 +26,20 @@ class Accumulator():
         self.dict = dict()
 
 
+class InfiniteDataLoader(DataLoader):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Initialize an iterator over the dataset.
+        self.dataset_iterator = super().__iter__()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            batch = next(self.dataset_iterator)
+        except StopIteration:
+            # Dataset exhausted, use a new fresh iterator.
+            self.dataset_iterator = super().__iter__()
+            batch = next(self.dataset_iterator)
+        return batch