Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| import os | |
| import shutil | |
| import torch | |
| def get_checkpoint_files(s): | |
| s = s.strip() | |
| if ',' in s: | |
| return [get_checkpoint_files(chunk) for chunk in s.split(',')] | |
| return 'last.ckpt' if s == 'last' else f'{s}.ckpt' | |
| def main(args): | |
| checkpoint_fnames = get_checkpoint_files(args.epochs) | |
| if isinstance(checkpoint_fnames, str): | |
| checkpoint_fnames = [checkpoint_fnames] | |
| assert len(checkpoint_fnames) >= 1 | |
| checkpoint_path = os.path.join(args.indir, 'models', checkpoint_fnames[0]) | |
| checkpoint = torch.load(checkpoint_path, map_location='cpu') | |
| del checkpoint['optimizer_states'] | |
| if len(checkpoint_fnames) > 1: | |
| for fname in checkpoint_fnames[1:]: | |
| print('sum', fname) | |
| sum_tensors_cnt = 0 | |
| other_cp = torch.load(os.path.join(args.indir, 'models', fname), map_location='cpu') | |
| for k in checkpoint['state_dict'].keys(): | |
| if checkpoint['state_dict'][k].dtype is torch.float: | |
| checkpoint['state_dict'][k].data.add_(other_cp['state_dict'][k].data) | |
| sum_tensors_cnt += 1 | |
| print('summed', sum_tensors_cnt, 'tensors') | |
| for k in checkpoint['state_dict'].keys(): | |
| if checkpoint['state_dict'][k].dtype is torch.float: | |
| checkpoint['state_dict'][k].data.mul_(1 / float(len(checkpoint_fnames))) | |
| state_dict = checkpoint['state_dict'] | |
| if not args.leave_discriminators: | |
| for k in list(state_dict.keys()): | |
| if k.startswith('discriminator.'): | |
| del state_dict[k] | |
| if not args.leave_losses: | |
| for k in list(state_dict.keys()): | |
| if k.startswith('loss_'): | |
| del state_dict[k] | |
| out_checkpoint_path = os.path.join(args.outdir, 'models', 'best.ckpt') | |
| os.makedirs(os.path.dirname(out_checkpoint_path), exist_ok=True) | |
| torch.save(checkpoint, out_checkpoint_path) | |
| shutil.copy2(os.path.join(args.indir, 'config.yaml'), | |
| os.path.join(args.outdir, 'config.yaml')) | |
| if __name__ == '__main__': | |
| import argparse | |
| aparser = argparse.ArgumentParser() | |
| aparser.add_argument('indir', | |
| help='Path to directory with output of training ' | |
| '(i.e. directory, which has samples, modules, config.yaml and train.log') | |
| aparser.add_argument('outdir', | |
| help='Where to put minimal checkpoint, which can be consumed by "bin/predict.py"') | |
| aparser.add_argument('--epochs', type=str, default='last', | |
| help='Which checkpoint to take. ' | |
| 'Can be "last" or integer - number of epoch') | |
| aparser.add_argument('--leave-discriminators', action='store_true', | |
| help='If enabled, the state of discriminators will not be removed from the checkpoint') | |
| aparser.add_argument('--leave-losses', action='store_true', | |
| help='If enabled, weights of nn-based losses (e.g. perceptual) will not be removed') | |
| main(aparser.parse_args()) | |