roedoejet
commented
2 months ago This is silly, but I tried to just move everything manually...it doesn't seem to work. training doesn't resume at the right epoch. There's probably something else that is mismatched but I can't figure out what.
import torch
from tqdm import tqdm
import regex as re
ev_ckpt = torch.load('/home/aip000/sgile/data/CRK/ev-sue-config/logs_and_checkpoints/VocoderExperiment/base/checkpoints/epoch=0-step=176.ckpt')
hfg_ckpt_d = torch.load('/home/aip000/tts/models/HiFi-GAN-TTS.pretrained/UNIVERSAL_V1/do_02500000')
hfg_ckpt_g = torch.load('/home/aip000/tts/models/HiFi-GAN-TTS.pretrained/UNIVERSAL_V1/g_02500000')
# Set optimizer states
ev_ckpt['optimizer_states'][0]['state'] = hfg_ckpt_d['optim_g']['state']
ev_ckpt['optimizer_states'][0]['param_groups'] = hfg_ckpt_d['optim_g']['param_groups']
ev_ckpt['optimizer_states'][1]['state'] = hfg_ckpt_d['optim_d']['state']
ev_ckpt['optimizer_states'][1]['param_groups'] = hfg_ckpt_d['optim_d']['param_groups']
# map generator/discriminator states
hfg_g_state_map = {k: f'generator.{k}' for k in hfg_ckpt_g['generator'].keys()} # key is old key, value is corresponding key in EV
hfg_msd_state_map = {k: f'msd.{k}' for k in hfg_ckpt_d['msd'].keys()} # key is old key, value is corresponding key in EV
hfg_mpd_state_map = {k: f'mpd.{k}' for k in hfg_ckpt_d['mpd'].keys()} # key is old key, value is corresponding key in EV
msd_pattern = re.compile(r'msd\.discriminators\.[1-9]')
mpd_pattern = re.compile(r'mpd\.discriminators')
resblock_pattern = re.compile(r'generator\.resblocks')
upsample_pattern = re.compile(r'generator\.ups')
convpre_pattern = re.compile(r'generator\.conv_p')
def add_parametrizations(pieces):
# add parametrizations
pieces.insert(-1, 'parametrizations')
# add weight
pieces.insert(-1, 'weight')
return pieces
def fix_weight_keys(map, patterns):
for k, v in tqdm(map.items()):
matches = [re.match(pattern, v) for pattern in patterns]
if any(matches):
pieces = v.split('.')
# convert weight_g/weight_v to original0/original1
if pieces[-1] == 'weight_g':
pieces[-1] = 'original0'
pieces = add_parametrizations(pieces)
elif pieces[-1] == 'weight_v':
pieces[-1] = 'original1'
pieces = add_parametrizations(pieces)
else:
continue
v = '.'.join(pieces)
map[k] = v
return map
hfg_g_state_map = fix_weight_keys(hfg_g_state_map, [resblock_pattern, upsample_pattern, convpre_pattern])
hfg_msd_state_map = fix_weight_keys(hfg_msd_state_map, [msd_pattern])
hfg_mpd_state_map = fix_weight_keys(hfg_mpd_state_map, [mpd_pattern])
gen_hfg_states = set(hfg_g_state_map.values())
msd_hfg_states = set(hfg_msd_state_map.values())
mpd_hfg_states = set(hfg_mpd_state_map.values())
gen_ev_states = set([x for x in ev_ckpt['state_dict'].keys() if x.startswith('generator')])
msd_ev_states = set([x for x in ev_ckpt['state_dict'].keys() if x.startswith('msd')])
mpd_ev_states = set([x for x in ev_ckpt['state_dict'].keys() if x.startswith('mpd')])
assert len(gen_ev_states) == len(gen_hfg_states) == 234
assert len(hfg_msd_state_map) == len(msd_ev_states) == 80
assert len(hfg_mpd_state_map) == len(mpd_ev_states) == 90
assert len(gen_ev_states - gen_hfg_states) == 0
assert len(msd_ev_states - msd_hfg_states) == 0
assert len(mpd_ev_states - mpd_hfg_states) == 0
# replace the weights
for k, v in hfg_g_state_map.items():
ev_ckpt['state_dict'][v] = hfg_ckpt_g['generator'][k]
for k, v in hfg_msd_state_map.items():
ev_ckpt['state_dict'][v] = hfg_ckpt_d['msd'][k]
for k, v in hfg_mpd_state_map.items():
ev_ckpt['state_dict'][v] = hfg_ckpt_d['mpd'][k]
# callbacks - meh? these don't seem like we need to move them
# Set steps
ev_ckpt['global_step'] = 5000000
# Set epoch
ev_ckpt['epoch'] = 12400
# loops (from EV universal)
ev_ckpt['loops'] = {'fit_loop': {'state_dict': {}, 'epoch_loop.state_dict': {'_batches_that_stepped': 2500000}, 'epoch_loop.batch_progress': {'total': {'ready': 2500000, 'completed': 2500000, 'started': 2500000, 'processed': 2500000}, 'current': {'ready': 12400, 'completed': 12400, 'started': 12400, 'processed': 12400}, 'is_last_batch': False}, 'epoch_loop.scheduler_progress': {'total': {'ready': 240, 'completed': 240}, 'current': {'ready': 0, 'completed': 0}}, 'epoch_loop.batch_loop.state_dict': {}, 'epoch_loop.batch_loop.optimizer_loop.state_dict': {}, 'epoch_loop.batch_loop.optimizer_loop.optim_progress': {'optimizer': {'step': {'total': {'ready': 5000000, 'completed': 5000000}, 'current': {'ready': 24800, 'completed': 24800}}, 'zero_grad': {'total': {'ready': 5000000, 'completed': 5000000, 'started': 5000000}, 'current': {'ready': 24800, 'completed': 24800, 'started': 24800}}}, 'optimizer_position': 2}, 'epoch_loop.batch_loop.manual_loop.state_dict': {}, 'epoch_loop.batch_loop.manual_loop.optim_step_progress': {'total': {'ready': 0, 'completed': 0}, 'current': {'ready': 0, 'completed': 0}}, 'epoch_loop.val_loop.state_dict': {}, 'epoch_loop.val_loop.dataloader_progress': {'total': {'ready': 120, 'completed': 120}, 'current': {'ready': 1, 'completed': 1}}, 'epoch_loop.val_loop.epoch_loop.state_dict': {}, 'epoch_loop.val_loop.epoch_loop.batch_progress': {'total': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}, 'current': {'ready': 36856, 'completed': 36856, 'started': 36856, 'processed': 36856}, 'is_last_batch': True}, 'epoch_progress': {'total': {'ready': 121, 'completed': 120, 'started': 121, 'processed': 121}, 'current': {'ready': 121, 'completed': 120, 'started': 121, 'processed': 121}}}, 'validate_loop': {'state_dict': {}, 'dataloader_progress': {'total': {'ready': 0, 'completed': 0}, 'current': {'ready': 0, 'completed': 0}}, 'epoch_loop.state_dict': {}, 'epoch_loop.batch_progress': {'total': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}, 'current': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}, 'is_last_batch': False}}, 'test_loop': {'state_dict': {}, 'dataloader_progress': {'total': {'ready': 0, 'completed': 0}, 'current': {'ready': 0, 'completed': 0}}, 'epoch_loop.state_dict': {}, 'epoch_loop.batch_progress': {'total': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}, 'current': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}, 'is_last_batch': False}}, 'predict_loop': {'state_dict': {}, 'dataloader_progress': {'total': {'ready': 0, 'completed': 0}, 'current': {'ready': 0, 'completed': 0}}, 'epoch_loop.state_dict': {}, 'epoch_loop.batch_progress': {'total': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}, 'current': {'ready': 0, 'completed': 0, 'started': 0, 'processed': 0}}}}
# LR Schedulers
ev_ckpt['hyper_parameters']['config']['training']['optimizer']['learning_rate'] = 0.0002
ev_ckpt['lr_schedulers'][0]['last_epoch'] = ev_ckpt['epoch']
ev_ckpt['lr_schedulers'][0]['_step_count'] = ev_ckpt['global_step']
torch.save(ev_ckpt, 'ev_hfg.ckpt')
SamuelLarkin
The Universal HiFiGAN checkpoint is very good, we could possibly just move the weights from it over to a shell EV model.