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