def load_wavernn(checkpoint_path: str) -> Tuple[WaveRNN, Dict[str, Any]]:
print(f'Loading voc checkpoint {checkpoint_path}')
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
config = checkpoint['config']
voc_model = WaveRNN.from_config(config)
voc_model.load_state_dict(checkpoint['model'])
print(f'Loaded model with step {voc_model.get_step()}')
return voc_model, config
parser.add_argument('--gta',
'-g',
action='store_true',
help='train wavernn on GTA features')
parser.add_argument('--config',
metavar='FILE',
default='config.yaml',
help='The config containing all hyperparams.')
args = parser.parse_args()
config = read_config(args.config)
paths = Paths(config['data_path'], config['voc_model_id'],
config['tts_model_id'])
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
print('Using device:', device)
print('\nInitialising Model...\n')
voc_model = WaveRNN.from_config(config).to(device)
dsp = DSP.from_config(config)
assert np.cumprod(
config['vocoder']['model']['upsample_factors'])[-1] == dsp.hop_length
optimizer = optim.Adam(voc_model.parameters())
restore_checkpoint(model=voc_model,
optim=optimizer,
path=paths.voc_checkpoints / 'latest_model.pt',
device=device)
voc_trainer = VocTrainer(paths=paths, dsp=dsp, config=config)
voc_trainer.train(voc_model, optimizer, train_gta=args.gta)
