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
train.py
+12 -11
View File
@@ -1,6 +1,7 @@
import os
import time
import argparse
import wandb
from cv2 import drawFrameAxes
import numpy as np
import tensorflow as tf
@@ -186,14 +187,14 @@ class Trainer(object):
return g_loss, g_gan_loss, pl_penalty
def train(self, dist_datasets, strategy):
# @tf.function
@tf.function
def dist_d_train_step(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_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
return mean_d_loss, mean_cond_loss
# @tf.function
@tf.function
def dist_d_train_step_reg(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)
@@ -202,13 +203,13 @@ class Trainer(object):
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
# @tf.function
@tf.function
def dist_g_train_step(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)
return mean_g_loss
# @tf.function
@tf.function
def dist_g_train_step_reg(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)
@@ -361,9 +362,9 @@ def main():
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('--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('--batch_size_per_replica', default=1, type=int)
parser.add_argument('--batch_size_per_replica', default=12, type=int)
args = vars(parser.parse_args())
# check tensorflow version
@@ -377,12 +378,12 @@ def main():
# network params
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
# featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
# featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
g_params = {
'z_dim': 16,
'w_dim': 16,
'z_dim': 512,
'w_dim': 512,
'labels_dim': 83,
'n_mapping': 8,
'resolutions': train_resolutions,
@@ -400,7 +401,7 @@ def main():
# prepare dataset
# 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.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
dataset = dataset.with_options(options)