import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import tensorflow as tf
import configs
import evaluation
from generating import inpainting, generate, k_nearest, intermediate
import trainv2 as train
import utils
import celeb_a_statistics
utils.manage_gpu_memory_usage()
EXPERIMENTS = {
"train": train.main,
"generate": generate.main,
"inpainting": inpainting.main,
"evaluation": evaluation.main,
"k_nearest": k_nearest.main,
"intermediate": intermediate.main,
"celeb_a_statistics": celeb_a_statistics.main
}
if __name__ == '__main__':
tf.get_logger().setLevel('ERROR')
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
args = utils.get_command_line_args()
configs.config_values = args
total_clip_norms = [
tf.cast(0, dtype=tf.float32),
tf.cast(0, dtype=tf.float32)
]
batched_dataset = dataset.batch(args.batch_size).prefetch(
buffer_size=tf.data.experimental.AUTOTUNE)
progress_bar = tqdm(batched_dataset, total=args.batches)
for index, data_batch in enumerate(progress_bar):
dnet_new_norm, snet_new_norm = train_step(
model, optimizer, data_batch, losses, clip_norms,
args.radious)
on_batch_end(epoch, index, dnet_new_norm, snet_new_norm,
total_clip_norms, losses, progress_bar,
args.batches)
validation_progress_bar = tqdm(
validation_dataset.batch(args.batch_size))
for validation_data_batch in validation_progress_bar:
validation_step(model, validation_data_batch, losses)
on_epoch_end(model, optimizer, epoch, losses, best_losses,
clip_norms, total_clip_norms,
args.checkpoint_directory)
if __name__ == '__main__':
manage_gpu_memory_usage()
train()