diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index 717a469..0000000 --- a/Dockerfile +++ /dev/null @@ -1,8 +0,0 @@ -FROM tensorflow/tensorflow:latest-gpu - -RUN apt-get update -RUN apt-get upgrade -y -RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip -y -RUN pip install -r requierements.txt - -# RUN wandb login \ No newline at end of file diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 1a59bcb..e591e8e 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -135,4 +135,16 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli zaxis = dict(nticks=4, range=[-1,1],),) ) - return fig \ No newline at end of file + return fig + + +def batch_to_dict(batch, signature): + result = dict() + for i, key in enumerate(signature.keys()): + result[key] = batch[i] + return result + +conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body'] +def nest_conditions(batch): + batch['conditions'] = {d: batch.pop(d) for d in conditions} + return batch diff --git a/dataset.py b/dataset.py index 55d4127..bcf57a7 100644 --- a/dataset.py +++ b/dataset.py @@ -6,6 +6,10 @@ import random import numpy as np import tensorflow as tf +from custom_utils.custom_utils import nest_conditions + + + class DatasetSimple(tf.data.Dataset): def _generator(folder): @@ -47,10 +51,182 @@ class DatasetSimple(tf.data.Dataset): args=(folder, ) ) -os.environ["CUDA_VISIBLE_DEVICES"] = "1" + + +class StaticDataset(): + def __init__(self, folder, signature) -> None: + self.store = dict() + self.signature = signature + self.folder = folder + self.possible_keys = ['uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body'] + self.keys = [value for value in signature.keys() if value in self.possible_keys] + + def read(self, id): + data = dict() + + if 'uv_static' in self.signature.keys(): + uv_static = np.load(os.path.join(self.folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static']=uv_static + if 'mesh_static' in self.signature.keys(): + mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static + if 'mesh_mask' in self.signature.keys(): + mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask + if 'uv_mask' in self.signature.keys(): + uv_mask = np.load(os.path.join(self.folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask'] = uv_mask + if 'uv_shaped_body' in self.signature.keys(): + uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body + + return data + + def __getitem__(self, id) -> dict: + # id = tf.strings.unicode_decode(id, 'UTF-8') + # id = tf.strings.as_string(id) + # id = id.numpy() + if id not in self.store: + self.store[id] = self.read(id) + + return self.store[id] + + def expand_batch(self, batch): + for key in self.keys: + batch[key] = tf.zeros(((0,) + self.signature[key])) + for i in range(batch['id'].shape[0]): + item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8')) + for key in self.keys: + batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0) + return batch + + + +class DatasetIdx(tf.data.Dataset): + + def __new__(cls, folder): + uv_static_folder = os.path.join(folder, 'pose_shape') + frames = os.listdir(uv_static_folder) + frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + random.shuffle(frames) + + return tf.data.Dataset.from_tensor_slices(frames) + + +# class DatasetFactory(tf.data.Dataset): +class DatasetFactory(): + + def __init__(self): + super(DatasetFactory, self).__init__() + + def _build_tf_signature(self, signature): + tf_signature = [] + for key in signature.keys(): + if key == 'id': + tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) + else: + tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) + return tf_signature + + + def get_dataset(self, folder, batch_size, signature): + # static_data = StaticDataset(folder, signature) + dataset_idx = DatasetIdx(folder) + + def read_item(frame): + frame_split = tf.strings.split(frame, sep='_') + id, frame_nr = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8") + # id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy() + # static_info = static_data[id] + item = dict() + + if 'id' in signature.keys(): + item['id'] = id + if 'pose_shape' in signature.keys(): + pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy')) + item['pose_shape'] = pose_shape + if 'uv_posed' in signature.keys(): + uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1]) + item['uv_posed'] = uv_posed + if 'uv_body_posed' in signature.keys(): + uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1]) + item['uv_body_posed'] = uv_body_posed + if 'mesh_body_posed' in signature.keys(): + mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) + item['mesh_body_posed'] = mesh_body_posed + if 'mesh_posed' in signature.keys(): + mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] + item['mesh_posed'] = mesh_posed + if 'mesh_unposed' in signature.keys(): + mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] + item['mesh_unposed'] = mesh_unposed + if 'uv_unposed' in signature.keys(): + 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]) + item['uv_unposed'] = uv_unposed + + + # item = { **static_info, **item } + as_list = [] + for key in signature.keys(): + if key in item: + as_list.append(item[key]) + + return as_list + + dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, + inp=[i], + Tout=self._build_tf_signature(signature) + ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat() + # dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat() + + return dataset + + # def __new__(cls, folder): + # return tf.data.Dataset.from_generator( + # cls._generator, + # # output_signature = (tf.TensorSpec(shape = (), dtype = tf.string), ), + # output_types=(tf.string), + # args=(folder, ) + # ) + + + +os.environ["CUDA_VISIBLE_DEVICES"] = "3" if __name__ == '__main__': + from custom_utils.custom_utils import batch_to_dict + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + # 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + # 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + # 'mesh_unposed': (14475, 3,), + # 'uv_unposed': (3, 512, 256,), + } + signature_static = { #'uv_static': (3, 512, 256,), + # 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + # 'uv_mask': (512, 256,), + # 'uv_shaped_body': (3, 512, 256,), + } + + dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train', + batch_size=1024, + signature=signature_dynamic) + static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train', + signature=signature_static) + + @tf.function + def euclid_error(mask, gt, pred): + return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))) / tf.reduce_sum(mask) + + all_error = [] + for batch in dataset: + batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) + err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], batch['mesh_body_posed']) + all_error.append(err.numpy()) + # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic))) + print(len(all_error)) + dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2) diff --git a/docker/Dockerfile b/docker/Dockerfile index 70b1b8a..a552b72 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,31 +1,10 @@ -FROM tensorflow/tensorflow:2.3.1-gpu -# FROM tensorflow/tensorflow:2.4.0rc0-gpu -MAINTAINER moono.song "toilety@gmail.com" -RUN apt-get update -y -RUN apt-get install -y build-essential -RUN pip install pillow +FROM tensorflow/tensorflow:2.9.0-gpu -# add sudoer -ARG USERNAME=anonymous -ARG USER_UID=1000 -ARG USER_GID=$USER_UID +RUN apt-get update +RUN apt-get upgrade -y +RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y +RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python -# Create the user -RUN groupadd --gid $USER_GID $USERNAME \ - && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \ - # - # [Optional] Add sudo support. Omit if you don't need to install software after connecting. - && apt-get update \ - && apt-get install -y sudo \ - && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \ - && chmod 0440 /etc/sudoers.d/$USERNAME -# ******************************************************** -# * Anything else you want to do like clean up goes here * -# ******************************************************** - -# [Optional] Set the default user. Omit if you want to keep the default as root. -USER $USERNAME - -# add port for multiworker training -EXPOSE 12345 \ No newline at end of file +WORKDIR /home/workspace/data +# RUN wandb login \ No newline at end of file diff --git a/requierements.txt b/docker/requierements.txt similarity index 63% rename from requierements.txt rename to docker/requierements.txt index 8574fc9..cf19a86 100644 --- a/requierements.txt +++ b/docker/requierements.txt @@ -5,4 +5,7 @@ wandb OpenEXR tensorflow-graphics scipy -plotly \ No newline at end of file +plotly +pandas +tensorboardX +opencv-python \ No newline at end of file diff --git a/load_models.py b/load_models.py index 6fdae31..19b4f4f 100644 --- a/load_models.py +++ b/load_models.py @@ -18,8 +18,9 @@ def load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=T } test_latent = tf.ones((1, g_params['z_dim']), dtype=tf.float32) - test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32)) + # test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32), tf.ones((1, 512, 256), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32)) # test_conditions = (tf.ones((1, 72), dtype=tf.float32), tf.ones((1, 11), dtype=tf.float32), tf.ones((1, 3, 512, 256), dtype=tf.float32)) + test_conditions = { 'pose_shape': tf.ones((1, 83), dtype=tf.float32), 'uv_static': tf.ones((1, 3, 512, 256), dtype=tf.float32), 'uv_mask': tf.ones((1, 512, 256), dtype=tf.float32), 'uv_shaped_body': tf.ones((1, 3, 512, 256), dtype=tf.float32) } # build generator model generator = Generator(g_params) diff --git a/losses.py b/losses.py index 7074949..8d84f7c 100644 --- a/losses.py +++ b/losses.py @@ -70,7 +70,10 @@ def g_logistic_ns_pathreg(real_images, generator, discriminator, z_dim, pl_minibatch = tf.maximum(1, tf.math.floordiv(batch_size, pl_minibatch_shrink)) pl_z = tf.random.normal(shape=[pl_minibatch, z_dim], dtype=tf.float32) - pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch]) + pl_conditions = { key: conditions[key][:pl_minibatch] for key in conditions.keys() } + # for key in conditions.keys(): + # pl_conditions[key] = conditions[key][:pl_minibatch] + # pl_conditions = (conditions[0][:pl_minibatch], conditions[1][:pl_minibatch], conditions[2][:pl_minibatch], conditions[3][:pl_minibatch], conditions[4][:pl_minibatch]) # forward pass fake_images = generator([z, conditions], ret_w_broadcasted=False, training=True) @@ -124,7 +127,7 @@ def collision_loss(V, B, mask, thr=.004): 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.reduce_sum((tf.minimum(dot - thr, 0) * mask) ** 2) / tf.reduce_sum(mask) + loss = tf.reduce_sum((tf.minimum(dot - thr, 0) * mask) ** 2, axis=-1) # / tf.reduce_sum(mask, axis=-1) return loss @@ -147,7 +150,7 @@ def edge_loss(x, y, mask): y_e = tf.sqrt(y_e + 1e-9) d_e = y_e - x_e d_e = d_e * e_mask - loss = tf.reduce_sum(d_e ** 2) / tf.reduce_sum(e_mask) + loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1) return loss diff --git a/stylegan2/generator.py b/stylegan2/generator.py index 2bcc193..96275b7 100644 --- a/stylegan2/generator.py +++ b/stylegan2/generator.py @@ -105,8 +105,8 @@ class Generator(tf.keras.Model): image_out = self.synthesis(w_broadcasted) # 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,:,:] + image_out = image_out + conditions['uv_shaped_body'] + image_out = image_out * conditions['uv_mask'][:,tf.newaxis,:,:] if ret_w_broadcasted: return image_out, w_broadcasted diff --git a/stylegan2/layers/embedding.py b/stylegan2/layers/embedding.py index 90e29c8..2242f93 100644 --- a/stylegan2/layers/embedding.py +++ b/stylegan2/layers/embedding.py @@ -97,6 +97,7 @@ class GarmentEmbedding(tf.keras.layers.Layer): def call(self, inputs, training=None, mask=None): x = inputs + x = tf.image.per_image_standardization(x) x = tf.transpose(x, [0,2,3,1]) x = self.encoder(x, training=training) diff --git a/stylegan2/layers/mapping_block.py b/stylegan2/layers/mapping_block.py index 430025d..4be814a 100644 --- a/stylegan2/layers/mapping_block.py +++ b/stylegan2/layers/mapping_block.py @@ -31,8 +31,8 @@ class Mapping(tf.keras.layers.Layer): x = latents # embed label if any - pose_shape_latent = self.pose_shape_embedding([conditions[0], conditions[1]]) - garment_latent = self.garment_embedding(conditions[2]) + pose_shape_latent = self.pose_shape_embedding([conditions['pose_shape'][:, :72], conditions['pose_shape'][:, 72:]]) + garment_latent = self.garment_embedding(conditions['uv_static']) x = tf.concat([x, pose_shape_latent, garment_latent], axis=1) # x = tf.concat([x, garment_latent], axis=1) diff --git a/train.py b/train.py index 3d12fb7..c907c3b 100644 --- a/train.py +++ b/train.py @@ -5,20 +5,20 @@ import wandb from cv2 import drawFrameAxes import numpy as np import tensorflow as tf -from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, plot_3d +from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, nest_conditions, plot_3d, batch_to_dict from tensorboardX import SummaryWriter 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 dataset import DatasetSimple, DatasetFactory, StaticDataset 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 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_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.') +run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_35', 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 and edge coeff. learns offsets over shaped body. fixed unposed uv map conversion. gnn discriminator. reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input') artifact = wandb.Artifact('code', type='code') artifact.add_file('train.py') artifact.add_file('losses.py') @@ -40,7 +40,7 @@ smpl = SMPL() class Trainer(object): def __init__(self, t_params, name): - self.cur_tf_ver = t_params['cur_tf_ver'] + # self.cur_tf_ver = t_params['cur_tf_ver'] self.use_tf_function = t_params['use_tf_function'] self.use_custom_cuda = t_params['use_custom_cuda'] self.model_base_dir = t_params['model_base_dir'] @@ -129,12 +129,13 @@ class Trainer(object): return params_copy def d_train_step(self, dist_inputs): - uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] - real_images = uv_unposed + # uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] + batch = dist_inputs[0] + real_images = batch['uv_unposed'] with tf.GradientTape() as d_tape: # compute losses - d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) + d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions']) # scale loss d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler @@ -146,13 +147,14 @@ class Trainer(object): return d_loss def d_train_step_reg(self, dist_inputs): - uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] - real_images = uv_unposed + # uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] + batch = dist_inputs[0] + real_images = batch['uv_unposed'] with tf.GradientTape() as d_tape: # compute losses d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator, - self.g_params['z_dim'], conditions=conditions) + self.g_params['z_dim'], conditions=batch['conditions']) r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval'] # scale losses @@ -168,12 +170,13 @@ class Trainer(object): def g_train_step(self, dist_inputs): - uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] - real_images = uv_unposed + # uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] + batch = dist_inputs[0] + real_images = batch['uv_unposed'] with tf.GradientTape() as g_tape: # compute losses - _, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) + _, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions']) # scale loss g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler @@ -183,14 +186,15 @@ class Trainer(object): return g_loss def g_train_step_reg(self, dist_inputs): - uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0] - real_images = uv_unposed + # uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0] + batch = dist_inputs[0] + real_images = batch['uv_unposed'] with tf.GradientTape() as g_tape: # compute losses _, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator, self.g_params['z_dim'], self.pl_mean, - self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions) + self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=batch['conditions']) pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval'] # scale loss @@ -206,28 +210,28 @@ class Trainer(object): def g_train_step_extra(self, dist_inputs): - uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0] - real_images = uv_unposed - uv_mask = conditions[3] + # uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0] + batch = dist_inputs[0] + real_images = batch['uv_unposed'] with tf.GradientTape() as g_tape: # compute losses - fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) + fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions']) mesh_pred_unposed = uv2mesh(fake_images) - mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed) + mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed) - e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask) - c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) - euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask) + e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask']) + c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']) + l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2]) + euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2) # scale loss g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler - # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler - # l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler + e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler + c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler + l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss @@ -238,32 +242,32 @@ class Trainer(object): def g_train_step_reg_extra(self, dist_inputs): # real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0] - uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0] - real_images = uv_unposed - uv_mask = conditions[3] + # uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0] + batch = dist_inputs[0] + real_images = batch['uv_unposed'] with tf.GradientTape() as g_tape: # compute losses fake_images, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator, self.g_params['z_dim'], self.pl_mean, - self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions) + self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=batch['conditions']) pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval'] mesh_pred_unposed = uv2mesh(fake_images) - mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed) + mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed) - e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask) - c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - l1_loss = 1e2 * tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) / tf.reduce_sum(uv_mask) - euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask) + e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask']) + c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']) + l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2]) + euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2) # scale loss pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler - # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler - # l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler + e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler + c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler + l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler # combine @@ -273,7 +277,7 @@ class Trainer(object): self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables)) return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim - def train(self, dist_datasets, strategy): + def train(self, dist_datasets: tf.data.Dataset, static_dataset: StaticDataset, dataset_signature: dict, strategy): @tf.function def dist_d_train_step(inputs): per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,)) @@ -306,9 +310,9 @@ class Trainer(object): def dist_g_train_step_extra(inputs): per_replica_losses = strategy.run(fn=self.g_train_step_extra, args=(inputs,)) mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) - mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[1], axis=None) - mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[2], axis=None) - mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None) + mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) + mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) + mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None) mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim @@ -319,9 +323,9 @@ class Trainer(object): mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) - mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None) - mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None) - mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[5], axis=None) + mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) + mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None) + mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[6], axis=None) mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[7], axis=None) return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim @@ -361,6 +365,7 @@ class Trainer(object): t_start = time.time() step = self.g_optimizer.iterations.numpy() for batch in dist_datasets: + batch = nest_conditions(static_dataset.expand_batch(batch_to_dict(batch, dataset_signature))) # d train step if (step + 1) % self.d_opt['reg_interval'] == 0: @@ -420,9 +425,6 @@ class Trainer(object): metric_euclid_error(euclid_error) metric_ms_ssim(ms_ssim) - # get current step - step += 1 - # save to tensorboard with train_summary_writer.as_default(): tf.summary.scalar('d_loss', metric_d_loss.result(), step=step) @@ -447,21 +449,18 @@ class Trainer(object): if step % self.image_summary_step == 0: # add summary image test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32) - test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32) + test_conditions = batch['conditions'] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32) summary_image = dist_gen_samples((test_z, test_conditions)) mesh_unposed = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):]) - mesh = smpl.run_garment(batch[1][0], batch[1][1], mesh_unposed) + mesh = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_unposed) # convert to tensor image summary_image = self.convert_per_replica_image(summary_image, strategy) with train_summary_writer.as_default(): tf.summary.image('images', summary_image, step=step) - template = uv2mesh(batch[1][2]) - # gt = uv2mesh(batch[0]) - plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch[6], template=template, body=batch[3], mask=batch[4].numpy(), step=step) - # plot_meshes(mesh_writer, garment=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step) + plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch['mesh_unposed'], template=batch['mesh_static'], body=batch['mesh_body_posed'], mask=batch['mesh_mask'].numpy(), step=step) # train_summary_writer.flush() @@ -479,6 +478,9 @@ class Trainer(object): if step >= self.max_steps: break + # get current step + step += 1 + # save last checkpoint step = self.g_optimizer.iterations.numpy() self.manager.save(checkpoint_number=step) @@ -492,7 +494,7 @@ class Trainer(object): fake_images_07 = self.g_clone([test_z, test_conditions], truncation_psi=0.7, training=False) # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res] - template = tf.concat([test_conditions[2], test_conditions[2]], axis=2) + template = tf.concat([test_conditions['uv_static'], test_conditions['uv_static']], axis=2) fake_image = tf.concat([fake_images_05, fake_images_07], axis=2) final_image = tf.concat([template, fake_image], axis=3) @@ -552,10 +554,11 @@ def main(): parser.add_argument('--train_res', default=512, type=int) parser.add_argument('--shuffle_buffer_size', default=1000, type=int) parser.add_argument('--batch_size_per_replica', default=4, type=int) + parser.add_argument('--name', type=str) args = vars(parser.parse_args()) # check tensorflow version - cur_tf_ver = check_tf_version() + # cur_tf_ver = check_tf_version() # GPU environment settings if args['allow_memory_growth']: @@ -594,7 +597,29 @@ 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, drop_remainder=True).repeat() + # dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size, drop_remainder=True).repeat() + + factory = DatasetFactory() + signature_dynamic = { 'id': (), + 'pose_shape': (72 + 11,), + # 'uv_posed': (3, 512, 256,), + # 'uv_body_posed': (3, 512, 256,), + 'mesh_body_posed': (14475, 3,), + 'mesh_posed': (14475, 3,), + 'mesh_unposed': (14475, 3,), + 'uv_unposed': (3, 512, 256,) } + signature_static = { 'uv_static': (3, 512, 256,), + 'mesh_static': (14475, 3,), + 'mesh_mask': (14475,), + 'uv_mask': (512, 256,), + 'uv_shaped_body': (3, 512, 256,), } + + dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train', + batch_size=4, + signature=signature_dynamic) + static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train', + signature=signature_static) + options = tf.data.Options() options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF @@ -607,7 +632,7 @@ def main(): # training parameters training_parameters = { # global params - 'cur_tf_ver': cur_tf_ver, + # 'cur_tf_ver': cur_tf_ver, 'use_tf_function': args['use_tf_function'], 'use_custom_cuda': args['use_custom_cuda'], 'model_base_dir': args['model_base_dir'], @@ -627,8 +652,8 @@ def main(): 'train_res': args['train_res'], } - trainer = Trainer(training_parameters, name=f'stylegan2-test-33-{args["train_res"]}x{args["train_res"]}') - trainer.train(dist_dataset, strategy) + trainer = Trainer(training_parameters, name=f'stylegan2-test-35-{args["train_res"]}x{args["train_res"]}') + trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy) return