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