def getTeacherRatio(epoch):
# from 1 to 0
return 1. - (1. / (epochs - 1)) * (epoch)
if __name__ == '__main__':
train_dataset = WordDataset('train')
test_dataset = WordDataset('test')
max_length = train_dataset.max_length
train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)
tense_list = test_dataloader.dataset.tense2idx.values()
model = CVAE(max_length)
model = model.to(device)
if args.load_model != None:
state_dict = torch.load(args.loadmodel)
model.load_state_dict(state_dict)
optimizer = optim.SGD(model.parameters(), lr=lr)
transformer = WordTransoformer()
trainset_size = len(train_dataloader.dataset)
testset_size = len(test_dataloader.dataset)
writer = SummaryWriter('logs/' + args.exp_name)
start = time.time()
best_bleu_score = 0
annealing_rate = 0.01
cycle = 15
# annealing_rate = 1./(args.warmup * len(train_dataloader.dataset))
print('Annealing: ', args.annealing)
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()
