added skinning and gnn discriminator

auto_encoder.py

@@ -1,17 +1,18 @@
 import os
 import time
 import argparse
+from cv2 import polarToCart
 import wandb
-from cv2 import drawFrameAxes
+import random
 import numpy as np
 import tensorflow as tf
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
 from tensorboardX import SummaryWriter
+import plotly.graph_objects as go
 
 from utils import str_to_bool
 from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
 from load_models import load_generator, load_discriminator
-from dataset import DatasetSimple
 from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
 # from losses import collision_penalty as collision_penalty
 from losses import collision_loss as collision_penalty
@@ -19,8 +20,8 @@ from convnext.models.convnext_tf import ConvNeXt
 from stylegan2_ae.generator import Generator
 
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
-run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_1_2', notes='convnext + stylegan generator')
+os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='ae_unposed_1', notes='convnext + stylegan generator')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('auto_encoder.py')
 artifact.add_file('stylegan2_ae/layers/mapping_block.py')
@@ -29,6 +30,35 @@ run.log_artifact(artifact)
 uv2mesh = UV2Mesh(res=256)
 
 
+class DatasetSimple(tf.data.Dataset):
+
+    def _generator(folder):
+        # Opening the file
+        folder = folder.decode("utf-8")
+        uv_folder = os.path.join(folder, 'uv_posed') 
+        frames = os.listdir(uv_folder)
+        random.shuffle(frames)
+
+        for frame in frames:
+            id, frame_nr = frame.split('.')[0].split('_')
+            mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
+            uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
+            mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
+            uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
+
+            yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
+
+    def __new__(cls, folder, res=256):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), 
+                                tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), 
+                                tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), 
+                                tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
+            args=(folder, )
+        )
+
+
 
 def filter_resolutions_featuremaps(resolutions, featuremaps, res):
     index = resolutions.index(res)
@@ -37,6 +67,43 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
     return filtered_resolutions, filtered_featuremaps
 
 
+def plot_3d(mesh_gt, mesh_pred, faces):
+    mesh_pred[:,1] += 2
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,1],
+            z=mesh_gt[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_pred[:,0],
+            y=mesh_pred[:,1],
+            z=mesh_pred[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig
+
+
 class AutoEncoder(tf.keras.Model):
     
     def __init__(self):
@@ -78,19 +145,18 @@ class AutoEncoder(tf.keras.Model):
 
 @tf.function
 def train_step(batch):
-    posed_uv, conditions,_, mesh_body_posed, mesh_mask = batch
-    uv_mask = conditions[3]
+    uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
 
     with tf.GradientTape() as tape:
 
-        pred = model([posed_uv, uv_mask])
+        pred = model([uv_unposed, uv_mask])
 
         mesh_pred = uv2mesh(pred)
-        mesh_orig = uv2mesh(posed_uv)
+        mesh_orig = mesh_unposed
 
         e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask)
         c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(posed_uv - pred)), axis=1) * uv_mask)
+        l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask)
         loss = l1_loss + e_loss + c_loss
 
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
@@ -102,16 +168,15 @@ def train_step(batch):
 
 @tf.function
 def valid_step(batch):
-    posed_uv, conditions,_, mesh_body_posed, mesh_mask = batch
-    uv_mask = conditions[3]
+    uv_unposed, mesh_unposed, uv_mask, mesh_mask = batch
 
-    pred = model([posed_uv, uv_mask])
+    pred = model([uv_unposed, uv_mask])
     mesh_pred = uv2mesh(pred)
-    mesh_orig = uv2mesh(posed_uv)
+    mesh_orig = mesh_unposed
 
     e_loss = 0.0#edge_loss(mesh_pred, mesh_orig, mesh_mask)
     c_loss = 0.0#collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
-    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(posed_uv - pred)), axis=1) * uv_mask)
+    l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(uv_unposed - pred)), axis=1) * uv_mask)
     loss = l1_loss + e_loss + c_loss
 
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
@@ -122,7 +187,7 @@ if __name__ == '__main__':
     allow_memory_growth()
     
     
-    name = 'test-9'
+    name = 'ae_unposed_1'
     dataset_train = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(8).repeat()
     dataset_valid = DatasetSimple('/home/workspace/data/dataset/stylegan_new/valid').cache().batch(8).repeat()
 
@@ -131,8 +196,7 @@ if __name__ == '__main__':
     optimizer = tf.optimizers.Adam()
     folder = 'models/ae_logs/'
     train_summary_writer = tf.summary.create_file_writer(folder + name + '/train')
-    valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
-    mesh_writer = SummaryWriter(folder + name + '/mesh')
+    # valid_summary_writer = tf.summary.create_file_writer(folder + name + '/valid')
 
     step = 0
     for batch in dataset_train:
@@ -142,18 +206,19 @@ if __name__ == '__main__':
             with train_summary_writer.as_default():
                 tf.summary.scalar(key, train_loss[key], step=step)
 
-        data_valid = iter(dataset_valid).get_next()
-        valid_loss = valid_step(data_valid)
+        data_valid = batch
+        # data_valid = iter(dataset_valid).get_next()
+        # valid_loss = valid_step(data_valid)
 
-        for key in valid_loss.keys():
-            with valid_summary_writer.as_default():
-                tf.summary.scalar(key, valid_loss[key], step=step)
+        # for key in valid_loss.keys():
+        #     with valid_summary_writer.as_default():
+        #         tf.summary.scalar(key, valid_loss[key], step=step)
 
         if step % 250 == 0:
-            pred = model([data_valid[0], data_valid[1][3]])
+            pred = model([data_valid[0], data_valid[2]])
             
             mesh_pred = uv2mesh(pred)
-            mesh_orig = uv2mesh(data_valid[0])
+            mesh_orig = data_valid[1]
             
             orig_img = tf.transpose(data_valid[0], [0,2,3,1])
             pred_img = tf.transpose(pred, [0,2,3,1])
@@ -162,12 +227,13 @@ if __name__ == '__main__':
                 tf.summary.image('images', img, step=step)
 
             F = np.load("custom_utils/smpl_faces.npy")
-            garment, faces = maskFacesBatch(mesh_pred.numpy(), F, data_valid[4].numpy().astype(bool))
-            gt, _ = maskFacesBatch(mesh_orig.numpy(), F, data_valid[4].numpy().astype(bool))
+            garment, faces = maskFacesBatch(mesh_pred.numpy(), F, data_valid[3].numpy().astype(bool))
+            gt = [mesh_orig[i].numpy()[data_valid[3][i].numpy().astype(bool), :] for i in range(pred.shape[0])]
             
-            for i in range(len(garment)):
-                mesh_writer.add_mesh('garment_pred_' + str(i), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
-                mesh_writer.add_mesh('garment_gt_' + str(i), vertices=gt[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
+            for i in range(pred.shape[0]):
+                fig = plot_3d(gt[i], garment[i], faces[i])
+                wandb.log({"garment_" + str(i): fig})
+                
 
         print(step)
 

custom_utils/graph_utils.py

@@ -0,0 +1,207 @@
+import scipy as sp
+import numpy as np
+import scipy.io as sio
+import tensorflow as tf
+from tensorflow.keras.activations import tanh, relu
+from tensorflow.keras.layers import Dense, BatchNormalization, Add, Flatten
+from custom_utils.custom_utils import UV2Mesh
+
+from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
+from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+
+
+def sparse_to_tensor(L):
+    idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
+    data = L.data.astype(dtype=np.float32)
+    shape = L.shape
+    return tf.SparseTensor(idx, data, shape)
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+
+
+def repeat_laplacian(L, times):
+    idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
+    idx_offset = np.linspace(0, L.shape[0]  * (times-1), times).astype(np.int32)
+    idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
+    idx += idx_offset[np.newaxis,:]
+    data = np.repeat(L.data, repeats=times, axis=0)
+    shape = (L.shape[0] * times, L.shape[1] * times)
+    return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
+
+
+
+class GraphConvolution(tf.keras.layers.Layer):
+    def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
+        super(GraphConvolution, self).__init__()
+        self.shape = shape
+        self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
+        self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
+        self.activation = activation or (lambda x: x)
+        self.normalization = normalization
+        self.batchSize = batchSize
+        self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
+        self.L = repeat_laplacian(L, self.batchSize)
+
+    def gather(self):
+        return [self.w, self.b]
+
+    def _single(self, X):
+        X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
+        X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
+        # Neigh conv
+        X1 = edge_convolution_template(
+            X1, sparse_to_tensor(self.L), sizes=None,
+            edge_function=lambda x, y: y, 
+            reduction='weighted', 
+            edge_function_kwargs={}
+        ) 
+        X = X0 + X1 + self.b
+        return X
+
+    def call(self, X, training):
+        X = merge_verts(X, self.batchSize)
+        X = self._single(X)
+        X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
+        X = self.activation(X)
+        if self.normalization:
+            X = self.normalization(X, training=training)
+        return X
+
+
+class GraphUnPool(tf.keras.layers.Layer):
+    def __init__(self, pool_map):
+        super(GraphUnPool, self).__init__()
+        self.pool_map = pool_map
+
+    def call(self, X):
+        # print(X.dtype)
+        # print(self.pool_map.dtype)
+        return unpool(X, self.pool_map, None)
+        # return tf.gather(X, self.pool_map, axis=1)
+
+
+
+class GraphPool(tf.keras.layers.Layer):
+    def __init__(self, pool_map, type, batch_size):
+        super(GraphPool, self).__init__()
+        self.type = type
+        self.batch_size = batch_size
+        self.indices = tf.cast(tf.linspace(0, batch_size * pool_map.shape[0], batch_size+1), dtype=tf.int32)
+        # self.pool_map = repeat_laplacian(pool_map, self.batch_size)
+        self.pool_map = np.squeeze(pool_map)
+
+    def _single(self, x):
+        return pool(x, sparse_to_tensor(self.pool_map), algorithm=self.type, sizes=None)
+
+    def call(self, x, training=None):
+        # x = merge_verts(x, self.batch_size)
+        # x = self._single(x)
+        # x = tf.stack([x[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
+
+        N, M = len(self.pool_map), len(set(self.pool_map)) # N = HQ verts, M = LQ verts
+        # Assign correspondences as V0_lq = [V0_hq, V1_hq, ..., VN_hq]
+        d = [[] for _ in range(M)] # list of lists
+        for i in range(N): d[self.pool_map[i]] += [i]
+        # max sublist len
+        m = max([len(d[i]) for i in range(M)])
+        # expand all sublist to max len
+        for i in range(M):
+            while len(d[i]) < m: 
+                # repeating elements will NOT change MIN/MAX output
+                d[i] += [d[i][0]] 
+ 
+        x = tf.reduce_max(tf.gather(x, d, axis=1), axis=2)
+        return x
+
+
+
+
+
+class GNNDiscriminator(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False, n_feature=256, pool_type="max", global_average=True):
+        super(GNNDiscriminator, self).__init__()
+
+        self.uv2mesh = UV2Mesh(res=256)        
+
+        laplacians = [sio.loadmat("custom_utils/laplacian/ssmpl" + ("_skirt" if isSkirt else "") + str(i) + ".mat")['L'].tocoo() for i in range(5)] 
+        poolings = [sio.loadmat("custom_utils/pooling/pool" + str(i) + str(i+1) + ".mat")['idx'] for i in range(5)]
+   
+        self.conv1 = GraphConvolution((3, int(n_feature/2**2)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
+        self.pool1 = GraphPool(pool_map=poolings[0], type=pool_type, batch_size=batchSize)
+
+        self.conv2 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**2)), activation=relu, L=laplacians[1], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_2')
+        self.pool2 = GraphPool(pool_map=poolings[1], type=pool_type, batch_size=batchSize)
+
+        self.custom_layers = []
+        for i in range(2,5):
+            self.custom_layers.append(GraphConvolution((int(n_feature/2**(4- (i if i == 2 else i-1) )), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
+            for j in range(2**(i-2)):
+                conv1 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
+                conv2 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
+                self.custom_layers.append([conv1, conv2])
+            self.custom_layers.append(GraphPool(pool_map=poolings[i], type=pool_type, batch_size=batchSize))
+            
+            self.custom_layers.append(Dense(256, relu))
+            self.custom_layers.append(BatchNormalization())
+            self.custom_layers.append(Dense(int(n_feature/2**(4-i)), relu))
+            self.custom_layers.append(BatchNormalization())
+
+        self.flatten = Flatten()
+        self.last = Dense(1)
+
+    def call(self, inputs, training=None):
+        inputs = self.uv2mesh(inputs)
+
+        x = self.conv1(inputs, training)
+        x = self.pool1(x)
+
+        x = self.conv2(x, training)
+        x = self.pool2(x)
+
+        for layer in self.custom_layers:
+            if isinstance(layer, list):
+                x1 = layer[0](x, training)
+                x2 = layer[1](x1, training)
+                x = Add()([x, x2])
+            else:
+                x = layer(x, training=training)
+
+        x = self.flatten(x)
+        x = self.last(x)
+
+        return x
+
+
+
+
+
+
+class DatasetTest(tf.data.Dataset):
+
+    def _generator():
+        for i in range(10):
+            yield np.array([i], dtype=np.int32)
+
+    def __new__(cls):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = tf.TensorSpec(shape = (1), dtype = tf.int32) ,
+            args=()
+        )
+
+
+
+if __name__ == '__main__':
+    import os
+    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+    batch_size = 4
+    input = tf.ones((batch_size, 14475, 3))
+
+    ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
+
+    x = ge(input)
+
+    print()

custom_utils/smpl.py

@@ -121,6 +121,7 @@ class SMPL:
 
     def run(self, R, betas, remove=True):
         batch_size = betas.shape[0]
+        # R = self.run_rodrigues(pose)
 
         #J_regressor = tf.gather(self.J_regressor, gender, axis=0)
         #shapedirs = tf.gather(self.shapedirs, gender, axis=0)

load_models.py

@@ -38,23 +38,20 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T
 
 
 def load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=True):
-    if custom_cuda:
-        from stylegan2.discriminator import Discriminator
-    else:
-        from stylegan2_ref.discriminator import Discriminator
+    # if custom_cuda:
+    #     from stylegan2.discriminator import Discriminator
+    # else:
+    #     from stylegan2_ref.discriminator import Discriminator
+
+    from custom_utils.graph_utils import GNNDiscriminator as Discriminator
 
-    if d_params is None:
-        d_params = {
-            'labels_dim': 0,
-            'resolutions': [4, 8, 16, 32, 64, 128, 256, 512, 1024],
-            'featuremaps': [512, 512, 512, 512, 512, 256, 128, 64, 32],
-        }
 
     res = d_params['resolutions'][-1]
-    test_images = tf.ones((1, 3, res, int(res/2)), dtype=tf.float32)
+    batch_size = d_params['batch_size']
+    test_images = tf.ones((batch_size, 3, res, int(res/2)), dtype=tf.float32)
 
     # build discriminator model
-    discriminator = Discriminator(d_params)
+    discriminator = Discriminator(batch_size, 32)
     _ = discriminator(test_images)
 
     if ckpt_dir is not None:

stylegan2/generator.py

@@ -17,6 +17,8 @@ class Generator(tf.keras.Model):
         self.featuremaps = g_params['featuremaps']
         self.w_ema_decay = 0.995
         self.style_mixing_prob = 0.9
+        self.zero_body = tf.ones((4,3,512,256)) * tf.transpose(g_params['zero_body'], [2,0,1])[tf.newaxis,:,:,:]
+
 
         self.n_broadcast = len(self.resolutions) * 2
         self.mixing_layer_indices = tf.range(self.n_broadcast, dtype=tf.int32)[tf.newaxis, :, tf.newaxis]
@@ -102,7 +104,8 @@ class Generator(tf.keras.Model):
             w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff)
 
         image_out = self.synthesis(w_broadcasted)
-        # image_out = image_out + conditions[4]
+        # image_out = image_out + conditions[2]
+        image_out = image_out + self.zero_body[:image_out.shape[0]]
         image_out = image_out * conditions[3][:,tf.newaxis,:,:]
 
         if ret_w_broadcasted:

train.py

@@ -18,7 +18,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.smpl import SMPL
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_28', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff')
+run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_33', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff. learns offsets over zero body. fixed unposed uv map conversion. gnn discriminator.')
 artifact = wandb.Artifact('code', type='code')
 artifact.add_file('train.py')
 artifact.add_file('losses.py')
@@ -579,8 +579,10 @@ def main():
         'n_mapping': 8,
         'resolutions': train_resolutions,
         'featuremaps': train_g_featuremaps,
+        'zero_body': np.load('custom_utils/uv_zero_body.npy')
     }
     d_params = {
+        'batch_size': args['batch_size_per_replica'],
         'labels_dim': -1,
         'resolutions': train_resolutions,
         'featuremaps': train_d_featuremaps,
@@ -592,7 +594,7 @@ def main():
 
     # prepare dataset
     # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
-    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size).repeat()
+    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size, drop_remainder=True).repeat()
 
     options = tf.data.Options()
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
@@ -625,7 +627,7 @@ def main():
             'train_res': args['train_res'],
         }
 
-        trainer = Trainer(training_parameters, name=f'stylegan2-test-28-{args["train_res"]}x{args["train_res"]}')
+        trainer = Trainer(training_parameters, name=f'stylegan2-test-33-{args["train_res"]}x{args["train_res"]}')
         trainer.train(dist_dataset, strategy)
     return