static encoder runs, but doesn't work yet

configs/ae_static.py

@@ -0,0 +1,71 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'static_torch_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'], 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

@@ -1,12 +1,16 @@
 import json
 import numpy as np
-import tensorflow as tf
+# import tensorflow as tf
+import torch
+import torch.nn as nn
 import scipy as sp
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 
-from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
+# from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
-from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
 
 
 def m2uv_array(m2uv, is_smpl_orig=False):
@@ -23,8 +27,8 @@ def m2uv_array(m2uv, is_smpl_orig=False):
     return result
 
 
-class UV2Mesh(tf.keras.Model):
+class UV2Mesh(nn.Module):
-    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False, is_adaptive=False):
+    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
         super(UV2Mesh, self).__init__()
         self.res = res
         uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy')
@@ -33,48 +37,40 @@ class UV2Mesh(tf.keras.Model):
         m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig)
         self.uv_pixels = self.uv_map_discreet[m2uv]
         self.smpl_res = 14475 if not is_smpl_orig else 3638
-        self.is_adaptive = is_adaptive
-
-        if self.is_adaptive:
-            laplacian = sp.sparse.load_npz("custom_utils/laplacian/laplacian0" + ("_skirt" if is_skirt else "") + ".npz")
-            self.gc = GraphConvolution((3, 3), L=laplacian, batchSize=4) 
 
 
-    @tf.function
     def uv_to_mesh(self, uv):
         if uv.shape[1] == 3:
-            uv = tf.transpose(uv, [0,2,3,1])
+            uv = torch.permute(uv, [0,2,3,1])
         if uv.shape[1] != 512:
-            uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5)
+            uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
             
-        batch_size = uv.shape[0]
+        # batch_size = uv.shape[0]
 
-        x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
+        # z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
-        x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
+        # y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1)
-        z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
+        # z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1)
 
 
-        x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
+        # x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
-        y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
+        # y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
-        z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
+        # z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
-        x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
+        # x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
-        y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
+        # y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
-        z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
+        # z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)
 
-        p1 = tf.stack((x1, y1, z1), axis=-1)
+        # p1 = tf.stack((x1, y1, z1), axis=-1)
-        p2 = tf.stack((x2, y2, z2), axis=-1)
+        # p2 = tf.stack((x2, y2, z2), axis=-1)
 
-        mesh = (p1 + p2) / 2
+        mesh = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]], dim=2)
         return mesh
 
 
     def __call__(self, uv):
         x = self.uv_to_mesh(uv)
-        if self.is_adaptive:
-            x = self.gc(x)
         return x
 
 
@@ -313,10 +309,16 @@ def load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False
 
 
 def get_summary_str(model):
-    lines = []
+    # from torchview import draw_graph
-    model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
+    # draw_graph(model, input_size=((4,3,512,256), (4,512,256)), expand_nested=True)
+
+    # lines = []
+    # model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True)
     # Add initial spaces to avoid markdown formatting in TensorBoard
-    return lines
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(param.numel() for param in model.parameters())
+
+    return [f'all_params: {total_params}, trainable: {trainable_params}']
 
 def get_summary_html(model):
     lines = get_summary_str(model)
@@ -341,7 +343,7 @@ def merge_verts(Vs, batchDim):
     return V
 
 
-class GraphConvolution(tf.keras.layers.Layer):
+class GraphConvolution(nn.Module):
     def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
         super(GraphConvolution, self).__init__()
         self.shape = shape
@@ -405,4 +407,46 @@ def group_inputs_for_levels(batch, config):
             uv_orig = batch['uv_posed' + res_modifier]
         levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']])
     conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']]
-    return levels_input, conds
+    return levels_input, conds
+
+
+class Accumulator():
+    def __init__(self):
+        self.dict = dict()
+
+    def update(self, dict):
+        for key in dict.keys():
+            if key in self.dict.keys():
+                self.dict[key].append(dict[key])
+            else:
+                self.dict[key] = [dict[key]]
+
+    def get(self):
+        result = dict()
+        for key in self.dict.keys():
+            result[key] = np.mean(self.dict[key])
+        return result
+
+    def reset(self):
+        self.dict = dict()
+
+
+from stl import mesh as mesh_stl
+
+def save_stl(file, mesh, faces):
+    obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype))
+    for i, f in enumerate(faces):
+        for j in range(3):
+            obj.vectors[i][j] = mesh[f[j],:]    
+
+    obj.save(file)
+
+
+def write_obj(file, vertices, faces, colors):
+    with open(file, 'w') as file:
+        for v, c in zip(vertices, colors):
+            line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n'
+            file.write(line)
+        for f in faces:
+            line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
+            file.write(line)

data/dataset.py

@@ -37,6 +37,8 @@ def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=
             frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
             frames = [frame.split('.')[0] + '_4' for frame in frames]
 
+        frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+
     if frames_file is not None:
         frames = []
         with open(os.path.join(folder, frames_file), 'r') as f:
@@ -122,7 +124,7 @@ class StaticCache():
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, signature, resolutions, frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
+    def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
         self.folder = folder
         self.device = device
         self.is_cape = is_cape
@@ -210,6 +212,10 @@ class Dataset(Dataset):
 
         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()}
 
+        for key, value in data.items():
+            if key.startswith('uv_') and value.shape[-1] == 3:
+                data[key] = torch.permute(value, [2,0,1])
+
         if self.cache is not None:
             self.cache[frame] = data
         return data

losses/losses.py

@@ -0,0 +1,244 @@
+import torch
+
+import numpy as np
+# import tensorflow_graphics as tfg
+import scipy as sp
+import scipy.io as sio
+
+from pytorch3d.structures import Meshes
+from scipy.spatial import cKDTree
+# from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
+# from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+
+
+
+# from hbertiche/PBNS
+""" COLLISION FUNCTIONS START """
+F_coll = np.load("smpl/smpl_faces.npy")
+F_coll_skirt = np.load("smpl/smpl_faces_skirt.npy")
+F_coll_orig = np.load("smpl/smpl_faces_orig.npy")
+F_norm = np.tile(F_coll[np.newaxis,:,:], [4,1,1])
+F_norm_skirt = np.tile(F_coll_skirt[np.newaxis,:,:], [4,1,1])
+F_norm_orig = np.tile(F_coll_orig[np.newaxis,:,:], [4,1,1])
+
+def _nearest_neighbour(V, B):
+	tree = cKDTree(B.numpy())
+	return tree.query(V.numpy(), workers=-1, distance_upper_bound=0.2)[1]
+
+# @tf.function
+# def _tf_nn_parallel(V, B):
+# 	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
+
+# @tf.function
+def collision_loss(V, B, mask, thr=.004):
+    mesh = Meshes(B, torch.tile(F_coll[None], [B.shape[0], 1, 1]))
+    B_vn = torch.stack(mesh.verts_normals_list())
+   
+    # compute correspondences
+    idx = torch.stack([_nearest_neighbour(V[i], B[i]) for i in range(V.shape[0])])
+    N, M = V.shape[:2]
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = torch.ones((N, M), dtype=torch.int32, device=V.get_device()) * np.array(list(range(N)), np.int32)[:,None]
+    idx = torch.stack((_idx, idx), axis=-1)
+    # compute loss
+    D = V - B[idx]
+    b_vn = B_vn[idx]
+    dot = torch.einsum('abc,abc->ab', D, b_vn)
+    thresholded = torch.minimum(dot - thr, 0)  * mask
+    loss = torch.sum(thresholded ** 2, axis=-1)
+    ratio = torch.sum((thresholded < 0).to(torch.float32)) / torch.sum(mask)
+
+    return loss, ratio
+
+
+# def collision_ratio(V, B, mask, thr=.004):
+    # B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+   
+    # # compute correspondences
+    # idx = _tf_nn_parallel(V, B)
+    # N, M = V.shape[:2]
+    # # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    # _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
+    # idx = tf.stack((_idx, idx), axis=-1)
+    # # compute loss
+    # D = V - tf.gather_nd(B, idx)
+    # b_vn = tf.gather_nd(B_vn, idx)
+    # dot = tf.einsum('abc,abc->ab', D, b_vn)
+    # loss = tf.minimum(dot - thr, 0)  * mask # / tf.reduce_sum(mask, axis=-1)
+    # ratio = tf.reduce_sum(tf.cast(loss < 0, dtype=tf.float32)) / tf.reduce_sum(mask)
+
+    # return ratio
+
+""" COLLISION FUNCTIONS END """
+
+
+laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
+edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
+
+laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
+edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
+
+
+def get_edge_mask(masks: torch.Tensor):
+    if edges.get_device() != masks.get_device():
+        edges.to(masks.get_device())
+    return torch.min(masks[:,edges], dim=-1)[0]
+
+def edge_loss(x, y, mask):
+    e_mask = get_edge_mask(mask)
+    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
+def get_edge_mask_skirt(masks: torch.Tensor):
+    if edges.get_device() != masks.get_device():
+        edges.to(masks.get_device())
+    return torch.min(masks[:,edges_skirt], dim=-1)[0]
+
+def edge_loss_skirt(x, y, mask):
+    e_mask = get_edge_mask_skirt(mask)
+    x_e = x[:,edges_skirt[:,0],:] - x[:,edges_skirt[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges_skirt[:,0],:] - y[:,edges_skirt[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    d_e = d_e * e_mask
+    loss = torch.sum(d_e ** 2, dim=-1)# / tf.reduce_sum(e_mask, axis=-1)
+    return loss
+
+
+# class PairwiseDistance(tf.keras.layers.Layer):
+#     def __init__(self, **kwargs):
+#         super(PairwiseDistance, self).__init__(**kwargs)
+
+#     def call(self, A, B):
+#         rA = tf.reduce_sum(tf.square(A), axis=-1)
+#         rB = tf.reduce_sum(tf.square(B), axis=-1)
+#         transpose_axes = [0, 2, 1] 
+#         distances = - 2*tf.matmul(A, tf.transpose(B, transpose_axes)) + rA[:, :, tf.newaxis] + rB[:, tf.newaxis, :]
+#         return distances
+
+
+# class NearestNeighbour(tf.keras.layers.Layer):
+#     def __init__(self, **kwargs):
+#         super(NearestNeighbour, self).__init__(**kwargs)
+
+#     def call(self, A, B):
+#         distances = PairwiseDistance(dtype=self.dtype)(A, B)
+#         nearest_neighbour = tf.argmin(distances, axis=-1)
+#         return tf.cast(nearest_neighbour, dtype=tf.int32)
+
+#
+# @tf.function
+def collision_penalty(va, vb, mask, eps=1e-3):
+    nb = vertex_normals(vb, tf.tile(F[None], [vb.shape[0], 1, 1]))
+    batch_size = tf.cast(tf.shape(va)[0], va.dtype)
+
+    closest_vertices = NearestNeighbour(dtype=va.dtype)(va, vb)
+    vb = tf.gather(vb, closest_vertices, batch_dims=1)
+    nb = tf.gather(nb, closest_vertices, batch_dims=1)
+
+    distance = tf.reduce_sum(nb*(va - vb), axis=-1) 
+    interpenetration = tf.maximum(eps - distance * mask, 0)
+
+    return tf.reduce_sum(interpenetration**3) / batch_size
+
+whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy')
+
+def uv_mask_loss(pred, uv_mask):
+    inv_mask = (1 - uv_mask)  * whole_garment_uv_mask
+    pred = tf.reduce_sum(pred, axis=1)
+    invalid = pred * inv_mask
+    invalid = invalid > 0.01
+    invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32)
+    total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32)
+    return invalid_count / total_count
+
+
+
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+    
+
+def _bend_loss_single(x, F, L, mask):
+    verts = tf.boolean_mask(x, mask, axis=0)
+    VN = vertex_normals(verts, F)
+    bend = edge_convolution_template(
+        VN, L, sizes=None,
+        edge_function=lambda x,y: (x - y)**2,
+        reduction='weighted',
+        edge_function_kwargs={}
+    )
+    loss = tf.reduce_sum(bend)
+    return loss
+
+def bend_loss(x, F, L, mask):
+    # x: predicted outfits
+    # F: faces
+    # L: laplacian
+    loss = tf.stack([_bend_loss_single(x[i], F[i], L[i], mask[i]) for i in range(x.shape[0])])
+    return loss
+
+
+# def bend_loss(x, F, L, mask):
+#     # x: predicted outfits
+#     # F: faces
+#     # L: laplacian
+#     # x = x * mask[:,:,tf.newaxis]
+#     batch_size = x.shape[0]
+#     verts, faces = maskFacesBatch(x, F, mask)
+#     nr_valid_verts = [tf.reduce_sum(mask[i]) for i in range(batch_size)]
+#     x = tf.concat(verts, axis=0)
+#     F = tf.concat([faces[i] + tf.reduce_sum(nr_valid_verts[:i]) for i in range(1, batch_size)])
+#     F = tf.concat((faces[0], F), axis=0)
+#     VN = vertex_normals(x, F)
+#     bend = edge_convolution_template(
+#         VN, L, sizes=None,
+#         edge_function=lambda x,y: (x - y)**2,
+#         reduction='weighted',
+#         edge_function_kwargs={}
+#     )
+#     loss = tf.reduce_sum(bend, axis=[1])
+#     loss = tf.reshape(loss, (batch_size, 14475))
+#     loss = tf.reduce_mean(loss, axis=1)
+#     return loss
+
+
+def log_normal_pdf(sample, mean, logvar, raxis=1):
+    log2pi = tf.math.log(2. * np.pi)
+    res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis)
+    return res
+
+
+def kl_loss(mean, logvar):
+    kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
+    kl_loss = tf.reduce_mean(kl_loss)
+    kl_loss *= -0.5
+    return kl_loss
+
+
+def normal_loss(pred, gt, mask, is_skirt=False, is_smpl_orig=False):
+    f_norm = None
+    if not is_skirt and not is_smpl_orig:
+        f_norm = F_norm
+    elif is_skirt:
+        f_norm = F_norm_skirt
+    elif is_smpl_orig:
+        f_norm = F_norm_orig
+    else:
+        raise NotImplementedError
+
+    pred_norm = vertex_normals(pred, f_norm)
+    gt_norm = vertex_normals(gt, f_norm)
+    return tf.reduce_mean(tf.reduce_mean((tf.reduce_sum(tf.math.abs(pred_norm - gt_norm), axis=-1) * mask), axis=-1))

models/base.py

@@ -1,134 +1,21 @@
-import tensorflow as tf
+import torch
-from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
+from torch import nn
+from models.convnext import ConvNeXt, CNBlock, CNBlockConfig
+from typing import Any, Callable, List, Optional, Sequence
 
-class Bottleneck(tf.keras.Model):
+
+class VariationalEncoder(nn.Module):
     
-    def __init__(self, dim, latent_dim, kernel_size=7):
+    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7):
-        super(Bottleneck, self).__init__()
-
-        self.bottleneck = tf.keras.Sequential(
-                [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-            
-        self.flatten = tf.keras.layers.Flatten()
-        self.dense = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.bottleneck(x, training=training)
-        x = self.flatten(x)
-        x = self.dense(x, training=training)
-
-        return x
-
-
-class VariationalEncoder(tf.keras.Model):
-    
-    def __init__(self, latent_dim, depths, dims, kernel_size=7):
         super(VariationalEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size)
 
-        self.flatten = tf.keras.layers.Flatten()
+        self.mean_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
-        self.mean = tf.keras.layers.Dense(latent_dim)
+        self.var_branch = nn.Sequential([CNBlock(block_setting[-1].out_channels, kernel_size=kernel_size) for _ in range(5)])
-        self.log_var = tf.keras.layers.Dense(latent_dim)
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.encoder(x, training=training)
-        x = self.flatten(x)
-        mean = self.mean(x)
-        log_var = self.log_var(x)
-
-        return mean, log_var
-
-
-class AEEncoder(tf.keras.Model):
-    
-    def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
-        super(AEEncoder, self).__init__()
-
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
-        self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
-
-
-    def call(self, inputs, training):
-        if len(inputs) == 2:
-            x, _ = inputs
-        else:
-            x = inputs
-
-        x = self.encoder(x, training=training)
-        x = self.bottleneck(x, training=training)
-
-        if len(inputs) == 2:
-            return x, inputs[1]
-        else:
-            return x
-
-
-class Decoder(tf.keras.Model):
-    
-    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
-        super(Decoder, self).__init__()
-
-        self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
-        self.reshape = tf.keras.layers.Reshape(input_size)
-        self.upscale = upscale
-        
-        if not is_label:
-            self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
-        else:
-            self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
-            
-
-    def call(self, inputs, training):
-        if len(inputs) == 2:
-            x, mask = inputs
-        else:
-            x = inputs
-
-        x = self.dense(x)
-        x = self.reshape(x)
-        x = self.decoder(x, training=training)
-
-        if self.upscale:
-            x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
-
-
-        if len(inputs) == 2:
-            if len(mask.shape) == 2:
-                x = x * mask[tf.newaxis,:,:,tf.newaxis]
-            else:
-                x = x * mask[:,:,:,tf.newaxis]
-
-        return x
-
-
-
-class VariationalEncoderNew(tf.keras.Model):
-    
-    def __init__(self, depths, dims, kernel_size=7):
-        super(VariationalEncoderNew, self).__init__()
-
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
-
-        self.mean_branch = tf.keras.Sequential(
-                [Block(dim=64, drop_path=0, kernel_size=kernel_size,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-        self.var_branch = tf.keras.Sequential(
-                [Block(dim=64, drop_path=0, kernel_size=kernel_size,
-                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
-            )
-
-        self.flatten_mean = tf.keras.layers.Flatten()
-        self.flatten_var = tf.keras.layers.Flatten()
 
+        self.flatten_mean = nn.Flatten()
+        self.flatten_var = nn.Flatten()
 
 
     def call(self, inputs, training):
@@ -144,23 +31,27 @@ class VariationalEncoderNew(tf.keras.Model):
         return mean, log_var
         
 
-
+class AEEncoder(nn.Module):
-class AEEncoderNew(tf.keras.Model):
     
-    def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
+    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'):
-        super(AEEncoderNew, self).__init__()
+        super(AEEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device)
-        self.act = tf.keras.layers.Activation(activation)
+        if activation == 'tanh':
+            self.act = nn.Tanh()
+        elif activation == 'sigmoid':
+            self.act = nn.Sigmoid()
+        else:
+            raise ValueError('Activation function not supported')
 
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         if len(inputs) == 2:
             x, _ = inputs
         else:
             x = inputs
 
-        x = self.encoder(x, training=training)
+        x = self.encoder(x)
         x = self.act(x)
 
         if len(inputs) == 2:
@@ -169,39 +60,36 @@ class AEEncoderNew(tf.keras.Model):
             return x
 
 
-class DecoderNew(tf.keras.Model):
+class Decoder(nn.Module):
     
-    def __init__(self, depths, dims, kernel_size=7, is_label=False):
+    def __init__(self, block_setting: CNBlockConfig, kernel_size=7, is_label=False, nr_of_classes=None, device='cuda'):
-        super(DecoderNew, self).__init__()
+        super(Decoder, self).__init__()
 
         self.is_label = is_label
 
-        if not is_label:
+        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.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
-        else:
-            self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
             
+        self.softmax = nn.Softmax(dim=1)
 
-    def call(self, inputs, training):
+    def forward(self, inputs):
         if len(inputs) == 2:
             x, mask = inputs
         else:
             x = inputs
 
-        x = self.decoder(x, training=training)
+        x = self.decoder(x)
 
         if self.is_label:
-            x = x * mask[:,:,:,tf.newaxis]
+            x = x * mask[:,:,:,None]
 
-            background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
+            background = torch.ones((mask.shape[0],512,256,1), device=mask.get_device()) * (1 - mask[:,:,:,None])
-            x = tf.concat((background, x), axis=-1)
+            x = torch.concat((background, x), axis=-1)
-
+            x = self.softmax(x)
-            x = tf.math.softmax(x)
 
         if len(inputs) == 2:
             if len(mask.shape) == 2:
-                x = x * mask[tf.newaxis,:,:,tf.newaxis]
+                x = x * mask[None,None,:,:]
             else:
-                x = x * mask[:,:,:,tf.newaxis]
+                x = x * mask[:,None,:,:]
 
         return x

models/cond_encoder.py

@@ -1,38 +1,21 @@
-import tensorflow as tf
+import torch
-from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
+from torch import nn
+from models.base import AEEncoder, Decoder
 
 
-class AutoEncoder(tf.keras.Model):
+class AutoEncoder(nn.Module):
     
     def __init__(self, config):
         super(AutoEncoder, self).__init__()
 
-        self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
+        self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], device=config['device'])
-        self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
+        self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'])
 
 
-    def call(self, inputs, training):
+    def forward(self, input):
-        x = inputs
+        x = input
 
-        x = self.encoder(x, training=training)
+        x = self.encoder(x)
-        x = self.decoder(x, training=training)
+        x = self.decoder(x)
-
-        return x
-
-
-class AutoEncoderNew(tf.keras.Model):
-    
-    def __init__(self, config):
-        super(AutoEncoderNew, self).__init__()
-
-        self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
-        self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
-
-
-    def call(self, inputs, training):
-        x = inputs
-
-        x = self.encoder(x, training=training)
-        x = self.decoder(x, training=training)
 
         return x

models/convnext.py

@@ -1,306 +1,424 @@
-from unittest import skip
+# taken from torchvision.models.convnext
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
-                                     GlobalAveragePooling2D, Layer,
-                                     LayerNormalization, Flatten, Reshape, UpSampling2D)
 
-def drop_path(inputs, drop_prob, is_training):
+from functools import partial
-    # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
+from typing import Any, Callable, List, Optional, Sequence
-    if (not is_training) or (drop_prob == 0.):
-        return inputs
 
-    # Compute keep_prob
+import torch
-    keep_prob = 1.0 - drop_prob
+from torch import nn, Tensor
+from torch.nn import functional as F
 
-    # Compute drop_connect tensor
+from torchvision.ops.misc import Conv2dNormActivation, Permute
-    random_tensor = keep_prob
+# from torchvision.ops.misc import ConvNormActivation as Conv2dNormActivation
-    shape = (tf.shape(inputs)[0],) + (1,) * \
+# from torchvision.ops import Permute
-        (len(tf.shape(inputs)) - 1)
+from torchvision.ops.stochastic_depth import StochasticDepth
-    random_tensor += tf.random.uniform(shape, dtype=inputs.dtype)
+from utils.convnext_utils import CNBlockConfig
-    binary_tensor = tf.floor(random_tensor)
-    output = tf.math.divide(inputs, keep_prob) * binary_tensor
-    return output
 
 
-class DropPath(tf.keras.layers.Layer):
+__all__ = [
-    # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
+    "ConvNeXt",
-    def __init__(self, drop_prob=None):
+    "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)
+        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__()
-        self.drop_prob = drop_prob
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
 
-    def call(self, x, training):
+        self.block = nn.Sequential(
-        return drop_path(x, self.drop_prob, training)
+            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 Block(Layer):
+class ConvNeXtBase(nn.Module):
-    r""" ConvNeXt Block. There are two equivalent implementations:
+    def __init__(
-    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+        self,
-    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+        block_setting: List[CNBlockConfig],
-    We use (2) as we find it slightly faster in PyTorch
+        stochastic_depth_prob: float = 0.0,
-
+        layer_scale: float = 1e-6,
-    Args:
+        block: Optional[Callable[..., nn.Module]] = None,
-        dim (int): Number of input channels.
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
-        drop_path (float): Stochastic depth rate. Default: 0.0
+        device: str = 'cuda',
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        kernel_size: int = 7,
-    """
+    ) -> None:
-
-    def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
         super().__init__()
-        self.dwconv = DepthwiseConv2D(
-            kernel_size=kernel_size, padding='same')  # depthwise conv
-        self.norm = LayerNormalization(epsilon=1e-6)
-        # pointwise/1x1 convs, implemented with linear layers
-        self.pwconv1 = Dense(4 * dim)
-        self.act = tf.keras.activations.gelu
-        self.pwconv2 = Dense(dim)
-        self.drop_path = DropPath(drop_path)
-        self.dim = dim
-        self.layer_scale_init_value = layer_scale_init_value
-        self.prefix = prefix
 
-    def build(self, input_shape):
+        if not block_setting:
-        self.gamma = tf.Variable(
+            raise ValueError("The block_setting should not be empty")
-            initial_value=self.layer_scale_init_value * tf.ones((self.dim)),
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
-            trainable=True,
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
-            name=f'{self.prefix}/gamma')
-        self.built = True
 
-    def call(self, x, training):
+        if block is None:
-        input = x
+            block = CNBlock
-        x = self.dwconv(x, training=training)
+
-        # x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        if norm_layer is None:
-        x = self.norm(x, training=training)
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
-        x = self.pwconv1(x, training=training)
+
+        # 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] = []
+
+        # 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)
+        
+
+        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
+            self.stages.append(nn.Sequential(*stage).to(device))
+            if cnf.out_channels is not None:
+                # Downsampling
+                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.act = nn.Tanh()
+
+        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.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.pwconv2(x, training=training)
-        if self.gamma is not None:
-            x = self.gamma * x
-        # x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
 
-        x = input + self.drop_path(x, training=training)
+        # x = self.permute_cl(x)
         return x
 
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
 
-class ConvNeXt(tf.keras.Model):
-    r""" ConvNeXt
-        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
 
-    Args:
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        include_top (bool): whether to add head or just use it as feature extractor. Default: True
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
 
-    def __init__(self, num_classes=1000, kernel_size=7,
+class ConvNeXtReverseBase2(nn.Module):
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
+    def __init__(
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+        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.include_top = include_top
-        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        # stem = tf.keras.Sequential([
-        #     Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
-        #     LayerNormalization(epsilon=1e-6)]
-        # )
-        stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
-        self.downsample_layers.append(stem)
-        for i in range(len(depths) - 1):
-            # downsample_layer = tf.keras.Sequential([
-            #     LayerNormalization(epsilon=1e-6),
-            #     Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            # )
-            downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            self.downsample_layers.append(downsample_layer)
 
-        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        self.use_skip = use_skip
-        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(len(depths)):
-            stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
-                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
 
-        if self.include_top:
+        if not block_setting:
-            self.avg = GlobalAveragePooling2D()
+            raise ValueError("The block_setting should not be empty")
-            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
-            self.head = Dense(num_classes)
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
-        else:
-            self.avg = None
-            self.norm = None
-            self.head = None
 
-    def forward_features(self, x, training):
+        if block is None:
-        for i in range(len(self.downsample_layers)):
+            block = CNBlock
-            # x = self.downsample_layers[i](x)
+
-            x = self.downsample_layers[i][0](x, training=training)
+        if norm_layer is None:
-            x = self.downsample_layers[i][1](x, training=training)
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
-            x = self.stages[i](x, training=training)
+
+        # 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),
+                    ).to(device)
+
+
+        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] = []
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # 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
+                self.stages.append(nn.Sequential(*stage).to(device))
+            else:
+                self.stages.append(nn.Identity().to(device))
+
+            if cnf.out_channels is not None:
+                # Upsampling
+                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.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)
+
+
+        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:
+        # 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.last(x)
+        x = self.activation(x)
+        # x = self.permute_cl(x)
+        
         return x
 
-    def call(self, x, training):
+    def forward(self, x) -> Tensor:
-        x = self.forward_features(x, training=training)
+        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,
+    reversed: bool,
+    is_label: bool,
+    nr_of_labels: Optional[int],
+    kernel_size: int,
+    device='cuda',
+    **kwargs: Any,
+) -> 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)
+    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
         
 
-class ConvNeXtReverse(tf.keras.Model):
-    r""" ConvNeXt
-        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
 
-    Args:
+if __name__ == '__main__':
-        num_classes (int): Number of classes for classification head. Default: 1000
+    model = ConvNeXt().to('cuda')
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+    test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        include_top (bool): whether to add head or just use it as feature extractor. Default: True
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
-
-    def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
-                 ):
-        super().__init__()
-        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        stem = tf.keras.Sequential([
-            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
-            # UpSampling2D((4, 4), interpolation='nearest'),
-            # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
-            # Dense(dims[0]),
-            LayerNormalization(epsilon=1e-6)]
-        )
-        self.upsample_layers.append(stem)
-        for i in range(len(depths) - 1):
-            downsample_layer = tf.keras.Sequential([
-                LayerNormalization(epsilon=1e-6),
-                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
-                # UpSampling2D((2,2), interpolation='nearest'),
-                # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
-                # Dense(dims[i+1]),
-            ])
-            self.upsample_layers.append(downsample_layer)
-
-        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
-        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(len(depths)):
-            stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
-                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
-
-        self.last = Dense(3, activation='tanh')
-
-    def call(self, x, training):
-        for i in reversed(range(len(self.upsample_layers))):
-            x = self.upsample_layers[i](x, training=training)
-            x = self.stages[i](x, training=training)
-        x = self.last(x, training=training)
-        return x
-
-
-class ConvNeXtReverseLabels(tf.keras.Model):
-    r""" ConvNeXt
-        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
-          https://arxiv.org/pdf/2201.03545.pdf
-
-    Args:
-        num_classes (int): Number of classes for classification head. Default: 1000
-        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
-        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
-        include_top (bool): whether to add head or just use it as feature extractor. Default: True
-        drop_path_rate (float): Stochastic depth rate. Default: 0.
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
-    """
-
-    def __init__(self, 
-                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
-                 nr_of_classes=7):
-        super().__init__()
-        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
-        stem = tf.keras.Sequential([
-            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
-            LayerNormalization(epsilon=1e-6)]
-        )
-        self.downsample_layers.append(stem)
-        for i in range(len(depths) - 1):
-            downsample_layer = tf.keras.Sequential([
-                LayerNormalization(epsilon=1e-6),
-                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
-            )
-            self.downsample_layers.append(downsample_layer)
-
-        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
-        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
-        cur = 0
-        for i in range(len(depths)):
-            stage = tf.keras.Sequential(
-                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
-                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
-
-        # self.last = Dense(nr_of_classes, activation='softmax')
-        self.last = Dense(nr_of_classes - 1)
-
-    def call(self, x, training):
-        for i in reversed(range(len(self.downsample_layers))):
-            x = self.downsample_layers[i](x, training=training)
-            x = self.stages[i](x, training=training)
-        x = self.last(x, training=training)
-        return x
-
-
-
-
-
-model_urls = {
-    "convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5",
-    "convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5",
-    "convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5",
-    "convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5",
-    "convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5",
-    "convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5",
-    "convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5",
-    "convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5",
-}
-
-
-model_configs = dict(
-    convnext_tiny=dict(
-        depths=[3, 3, 9, 3],
-        dims=[8, 16, 32, 512]
-    ),
-    convnext_small=dict(
-        depths=[3, 3, 27, 3],
-        dims=[96, 192, 384, 768]
-    ),
-    convnext_base=dict(
-        depths=[3, 3, 27, 3],
-        dims=[128, 256, 512, 1024]
-    ),
-    convnext_large=dict(
-        depths=[3, 3, 27, 3],
-        dims=[192, 384, 768, 1536]
-    ),
-    convnext_xlarge=dict(
-        depths=[3, 3, 27, 3],
-        dims=[256, 512, 1024, 2048]
-    ),
-)
 
+    pass

run.py

@@ -0,0 +1,9 @@
+import os
+from configs.ae_static import config
+
+
+if __name__ == '__main__':
+
+    name = config['name']
+    current_file = '/'.join(__file__.split('/')[-2:])
+    os.system(f'nohup python scripts/ae_static_train.py --name lhunor_{name} > nohup/torch/{name}.out 2>&1 &')

scripts/ae_static_train.py

@@ -0,0 +1,197 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from configs.ae_static import config
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+
+
+import time
+import argparse
+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
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch.optim as optim
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+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 models.cond_encoder import AutoEncoder
+from utils.checkpoints import CheckpointSaver
+from custom_utils.custom_utils import Accumulator
+
+
+def init_wandb():
+    run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp/wandb', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/ae_static_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+
+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)
+
+    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'])
+    
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
+        
+
+def valid_step(batch):
+    model.eval()
+
+    with torch.inference_mode():
+        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)
+
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+       
+
+def plot_everything(batch: dict, type: str):
+    pred = model([batch['uv_static'], batch['uv_mask']]).detach()
+    
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_static']
+    
+    orig_img = torch.permute(batch['uv_static'], [0,2,3,1])
+    pred_img = torch.permute(pred, [0,2,3,1])
+    img = torch.concat((orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
+
+    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+    garment_static, faces = maskFacesBatch(mesh_pred.detach().cpu().numpy(), F, mask)
+    gt = [mesh_orig[i].cpu().numpy()[mask[i], :] for i in range(pred.shape[0])]
+    
+    body = gt
+    garment_posed = garment_static
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_static[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    
+    return { **figs, **{ type + '/images': imgs } }
+
+
+def train():
+    valid_iterator = iter(valid_dataloader)
+
+    step = 0
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0:
+            ckpt_saver_enc(model.encoder, step, optimizer)
+            ckpt_saver_dec(model.decoder, step, optimizer)
+
+        print(step)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    init_wandb()
+    # device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--name', type=str)
+    args = vars(parser.parse_args())
+    
+    resolution_modifier = ''
+    if config['resolution'] != 512:
+        resolution_modifier = '_' + str(config['resolution'])
+    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)
+    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'])
+
+    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)
+
+    run_name = config['name']
+    print(f'started training: {run_name}')
+    train()
+
+
+

utils/convnext_utils.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+class CNBlockConfig:
+    # Stores information listed at Section 3 of the ConvNeXt paper
+    def __init__(
+        self,
+        input_channels: int,
+        out_channels: Optional[int],
+        num_layers: int,
+    ) -> None:
+        self.input_channels = input_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "input_channels={input_channels}"
+        s += ", out_channels={out_channels}"
+        s += ", num_layers={num_layers}"
+        s += ")"
+        return s.format(**self.__dict__)