mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-02-24 19:54:43 +08:00
80 lines
3.0 KiB
Python
Executable File
80 lines
3.0 KiB
Python
Executable File
#!/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())
|