| | """ |
| | 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: |
| | 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() |
| |
|
