Spaces:
Sleeping
Sleeping
File size: 2,504 Bytes
f4c3c2b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
"""
This file runs the main training/val loop
"""
import os
import json
import math
import sys
import pprint
import torch
from argparse import Namespace
sys.path.append(".")
sys.path.append("..")
from options.train_options import TrainOptions
from training.coach import Coach
def main():
opts = TrainOptions().parse()
previous_train_ckpt = None
if opts.resume_training_from_ckpt:
opts, previous_train_ckpt = load_train_checkpoint(opts)
else:
setup_progressive_steps(opts)
create_initial_experiment_dir(opts)
coach = Coach(opts, previous_train_ckpt)
coach.train()
def load_train_checkpoint(opts):
train_ckpt_path = opts.resume_training_from_ckpt
previous_train_ckpt = torch.load(opts.resume_training_from_ckpt, map_location='cpu')
new_opts_dict = vars(opts)
opts = previous_train_ckpt['opts']
opts['resume_training_from_ckpt'] = train_ckpt_path
update_new_configs(opts, new_opts_dict)
pprint.pprint(opts)
opts = Namespace(**opts)
if opts.sub_exp_dir is not None:
sub_exp_dir = opts.sub_exp_dir
opts.exp_dir = os.path.join(opts.exp_dir, sub_exp_dir)
create_initial_experiment_dir(opts)
return opts, previous_train_ckpt
def setup_progressive_steps(opts):
log_size = int(math.log(opts.stylegan_size, 2))
num_style_layers = 2*log_size - 2
num_deltas = num_style_layers - 1
if opts.progressive_start is not None: # If progressive delta training
opts.progressive_steps = [0]
next_progressive_step = opts.progressive_start
for i in range(num_deltas):
opts.progressive_steps.append(next_progressive_step)
next_progressive_step += opts.progressive_step_every
assert opts.progressive_steps is None or is_valid_progressive_steps(opts, num_style_layers), \
"Invalid progressive training input"
def is_valid_progressive_steps(opts, num_style_layers):
return len(opts.progressive_steps) == num_style_layers and opts.progressive_steps[0] == 0
def create_initial_experiment_dir(opts):
if os.path.exists(opts.exp_dir):
raise Exception('Oops... {} already exists'.format(opts.exp_dir))
os.makedirs(opts.exp_dir)
opts_dict = vars(opts)
pprint.pprint(opts_dict)
with open(os.path.join(opts.exp_dir, 'opt.json'), 'w') as f:
json.dump(opts_dict, f, indent=4, sort_keys=True)
def update_new_configs(ckpt_opts, new_opts):
for k, v in new_opts.items():
if k not in ckpt_opts:
ckpt_opts[k] = v
if new_opts['update_param_list']:
for param in new_opts['update_param_list']:
ckpt_opts[param] = new_opts[param]
if __name__ == '__main__':
main()
|