CVPR/lama-example

1#!/usr/bin/env python32 3import os4import shutil5 6import torch7 8 9def get_checkpoint_files(s):10    s = s.strip()11    if ',' in s:12        return [get_checkpoint_files(chunk) for chunk in s.split(',')]13    return 'last.ckpt' if s == 'last' else f'{s}.ckpt'14 15 16def main(args):17    checkpoint_fnames = get_checkpoint_files(args.epochs)18    if isinstance(checkpoint_fnames, str):19        checkpoint_fnames = [checkpoint_fnames]20    assert len(checkpoint_fnames) >= 121 22    checkpoint_path = os.path.join(args.indir, 'models', checkpoint_fnames[0])23    checkpoint = torch.load(checkpoint_path, map_location='cpu')24    del checkpoint['optimizer_states']25 26    if len(checkpoint_fnames) > 1:27        for fname in checkpoint_fnames[1:]:28            print('sum', fname)29            sum_tensors_cnt = 030            other_cp = torch.load(os.path.join(args.indir, 'models', fname), map_location='cpu')31            for k in checkpoint['state_dict'].keys():32                if checkpoint['state_dict'][k].dtype is torch.float:33                    checkpoint['state_dict'][k].data.add_(other_cp['state_dict'][k].data)34                    sum_tensors_cnt += 135            print('summed', sum_tensors_cnt, 'tensors')36 37        for k in checkpoint['state_dict'].keys():38            if checkpoint['state_dict'][k].dtype is torch.float:39                checkpoint['state_dict'][k].data.mul_(1 / float(len(checkpoint_fnames)))40 41    state_dict = checkpoint['state_dict']42 43    if not args.leave_discriminators:44        for k in list(state_dict.keys()):45            if k.startswith('discriminator.'):46                del state_dict[k]47 48    if not args.leave_losses:49        for k in list(state_dict.keys()):50            if k.startswith('loss_'):51                del state_dict[k]52 53    out_checkpoint_path = os.path.join(args.outdir, 'models', 'best.ckpt')54    os.makedirs(os.path.dirname(out_checkpoint_path), exist_ok=True)55 56    torch.save(checkpoint, out_checkpoint_path)57 58    shutil.copy2(os.path.join(args.indir, 'config.yaml'),59                 os.path.join(args.outdir, 'config.yaml'))60 61 62if __name__ == '__main__':63    import argparse64 65    aparser = argparse.ArgumentParser()66    aparser.add_argument('indir',67                         help='Path to directory with output of training '68                              '(i.e. directory, which has samples, modules, config.yaml and train.log')69    aparser.add_argument('outdir',70                         help='Where to put minimal checkpoint, which can be consumed by "bin/predict.py"')71    aparser.add_argument('--epochs', type=str, default='last',72                         help='Which checkpoint to take. '73                              'Can be "last" or integer - number of epoch')74    aparser.add_argument('--leave-discriminators', action='store_true',75                         help='If enabled, the state of discriminators will not be removed from the checkpoint')76    aparser.add_argument('--leave-losses', action='store_true',77                         help='If enabled, weights of nn-based losses (e.g. perceptual) will not be removed')78 79    main(aparser.parse_args())80