model, test_dataloader, tense_list) writer.add_scalar('Loss/reconstruction loss', rc_loss / trainset_size, epoch + 1) writer.add_scalar('Loss/KL loss', kl_loss / trainset_size, epoch + 1) writer.add_scalar('BLEU-4 score', average_bleu_score, epoch + 1) writer.add_scalar('Setting parameter/KL weight', kl_weight, epoch + 1) writer.add_scalar('Setting parameter/teacher forcing ratio', teacher_forcing_ratio, epoch + 1) writer.add_scalars( 'Comparison', { 'rc_loss': rc_loss / trainset_size, 'kl_loss': kl_loss / trainset_size, 'BLEU-4 socre': average_bleu_score, 'kl_weight': kl_weight, 'teacher_ratio': teacher_forcing_ratio, 'gaussian_score': gaussian_score }, epoch + 1) print('Epoch ', epoch + 1) print('Average Reconstruction loss: %f\n Average KL loss: %f' % (rc_loss / trainset_size, kl_loss / trainset_size)) print('Average BLEU-4 score: %f\n' % average_bleu_score) if average_bleu_score > best_bleu_score: record_score(average_bleu_score, gaussian_score, predict_list, generate_words, test_dataloader, transformer) best_bleu_score = average_bleu_score torch.save(model.state_dict(), 'model/checkpoint' + str(epoch) + '.pkl') end = time.time() print('Total training time: ' + str((end - start) // 60) + ' minutess')
def main(): parser = argparse.ArgumentParser( description='Train MVAE with VCC2018 dataset', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--train-dataset', help='Path of training dataset.', type=str, required=True) parser.add_argument('--val-dataset', help='Path of validation dataset.', type=str, required=True) parser.add_argument('--batch-size', '-b', help='Batch size.', type=int, default=32) parser.add_argument('--epochs', '-e', help='Number of epochs.', type=int, default=800) parser.add_argument('--eval-interval', help='Evaluate and save model every N epochs.', type=int, default=200, metavar='N') parser.add_argument('--gpu', '-g', help='GPU id. (Negative number indicates CPU)', type=int, default=-1) parser.add_argument('--learning-rate', '-l', help='Learning Rate.', type=float, default=1e-3) parser.add_argument('--output', help='Save model to PATH', type=str, default='./models') args = parser.parse_args() if not os.path.isdir(args.output): os.mkdir(args.output) if_use_cuda = torch.cuda.is_available() and args.gpu >= 0 if if_use_cuda: device = torch.device(f'cuda:{args.gpu}') cp.cuda.Device(args.gpu).use() else: device = torch.device('cpu') train_dataset = torch.load(args.train_dataset) train_dataloader = torch.utils.data.DataLoader( train_dataset, args.batch_size, shuffle=True) val_dataset = make_eval_set(args.val_dataset) baseline = baseline_ilrma(val_dataset, device) model = CVAE(n_speakers=train_dataset[0][1].size(0)).to(device) optimizer = torch.optim.Adam(model.parameters(), args.learning_rate) # TensorBoard writer = SummaryWriter() for epoch in range(1, args.epochs + 1): train(model, train_dataloader, optimizer, device, epoch, writer) if epoch % args.eval_interval == 0: validate(model, val_dataset, baseline, device, epoch, writer) # Save model model.cpu() path = os.path.join(args.output, f'model-{epoch}.pth') torch.save(model.state_dict(), path) model.to(device) writer.close()