luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

settings for the server

Browse Source
This commit is contained in:
luni
2022-05-25 20:00:12 +02:00
parent b098ecee9c
commit 564b9e96a7
2 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+1 -1
View File
@@ -17,7 +17,7 @@ class DatasetSimple(tf.data.Dataset):
for frame in frames: for frame in frames:
uv = np.load(os.path.join(uv_folder, frame)) uv = np.load(os.path.join(uv_folder, frame))
uv = uv['uv_gt'] uv = uv['uv_gt']
uv = cv2.resize(uv, [64, 128]) # uv = cv2.resize(uv, [64, 128])
uv = np.transpose(uv, [2, 0, 1]) uv = np.transpose(uv, [2, 0, 1])
pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy')) pose_shape = np.load(os.path.join(pose_hape_folder, frame.split('.')[0] + '.npy'))
train.py
+12 -11
View File
@@ -1,6 +1,7 @@
import os import os
import time import time
import argparse import argparse
import wandb
from cv2 import drawFrameAxes from cv2 import drawFrameAxes
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
@@ -186,14 +187,14 @@ class Trainer(object):
return g_loss, g_gan_loss, pl_penalty return g_loss, g_gan_loss, pl_penalty
def train(self, dist_datasets, strategy): def train(self, dist_datasets, strategy):
# @tf.function @tf.function
def dist_d_train_step(inputs): def dist_d_train_step(inputs):
per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,)) per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
return mean_d_loss, mean_cond_loss return mean_d_loss, mean_cond_loss
# @tf.function @tf.function
def dist_d_train_step_reg(inputs): def dist_d_train_step_reg(inputs):
per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,)) per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
@@ -202,13 +203,13 @@ class Trainer(object):
mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_cond_loss return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_cond_loss
# @tf.function @tf.function
def dist_g_train_step(inputs): def dist_g_train_step(inputs):
per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,)) per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None) mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
return mean_g_loss return mean_g_loss
# @tf.function @tf.function
def dist_g_train_step_reg(inputs): def dist_g_train_step_reg(inputs):
per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,)) per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,))
mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
@@ -361,9 +362,9 @@ def main():
parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=True) parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=True)
parser.add_argument('--model_base_dir', default='./models', type=str) parser.add_argument('--model_base_dir', default='./models', type=str)
parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str) parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
parser.add_argument('--train_res', default=128, type=int) parser.add_argument('--train_res', default=512, type=int)
parser.add_argument('--shuffle_buffer_size', default=1000, type=int) parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
parser.add_argument('--batch_size_per_replica', default=1, type=int) parser.add_argument('--batch_size_per_replica', default=12, type=int)
args = vars(parser.parse_args()) args = vars(parser.parse_args())
# check tensorflow version # check tensorflow version
@@ -377,12 +378,12 @@ def main():
# network params # network params
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
# featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4] # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res']) train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
g_params = { g_params = {
'z_dim': 16, 'z_dim': 512,
'w_dim': 16, 'w_dim': 512,
'labels_dim': 83, 'labels_dim': 83,
'n_mapping': 8, 'n_mapping': 8,
'resolutions': train_resolutions, 'resolutions': train_resolutions,
@@ -400,7 +401,7 @@ def main():
# prepare dataset # prepare dataset
# dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None) # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
dataset = DatasetSimple('/workspace/Dataset/stylegan').batch(global_batch_size)#.shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache() dataset = DatasetSimple('/home/worksapce/data/dataset/stylegan').batch(global_batch_size)#.shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache()
options = tf.data.Options() options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
dataset = dataset.with_options(options) dataset = dataset.with_options(options)