loss = train(epoch)
test_loss, test_f0_loss = test(epoch)
print(loss)
print(test_loss)
print('epoch [{}/{}], loss: {:.4f} test_loss: {:.4f}'.format(
epoch + 1, num_epochs, loss, test_loss))
# logging
loss_list.append(loss)
test_loss_list.append(test_loss)
test_f0_losses.append(test_f0_loss)
print(time.time() - start)
if args.model_path != '':
pre_trained_epoch = int(
args.model_path[args.model_path.index('model_') +
6:args.model_path.index('.pth')])
else:
pre_trained_epoch = 0
if epoch % 5 == 0:
torch.save(
model.state_dict(), args.output_dir + '/model_' +
str(epoch + pre_trained_epoch) + '.pth')
np.save(args.output_dir + '/loss_list.npy', np.array(loss_list))
np.save(args.output_dir + '/test_loss_list.npy', np.array(test_loss_list))
np.save(args.output_dir + '/test_f0loss_list.npy',
np.array(test_f0_losses))
def objective(trial):
mora_index_lists = sorted(
glob(join('data/basic5000/mora_index', "squeezed_*.csv")))
#mora_index_lists = mora_index_lists[:len(mora_index_lists)-5] # last 5 is real testset
mora_index_lists_for_model = [
np.loaadtxt(path).reshape(-1) for path in mora_index_lists
]
train_mora_index_lists = []
test_mora_index_lists = []
#train_files, test_files = train_test_split(files, test_size=test_size, random_state=random_state)
for i, mora_i in enumerate(mora_index_lists_for_model):
if (i - 1) % 20 == 0: #test
pass
elif i % 20 == 0: #valid
test_mora_index_lists.append(mora_i)
else:
train_mora_index_lists.append(mora_i)
num_lstm_layers = trial.suggest_int('num_lstm_layers', 1, 3)
z_dim = trial.suggest_categorical('z_dim', [
1,
2,
8,
])
num_class = trial.suggest_int('num_class', 2, 4)
model = VQVAE(num_class=num_class, num_layers=num_lstm_layers,
z_dim=z_dim).to(device)
optimizer = optim.Adam(model.parameters(), lr=2e-3) #1e-3
start = time.time()
# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function(recon_x, x, z, z_unquantized, beta=1):
MSE = F.mse_loss(
recon_x.view(-1), x.view(-1, ), reduction='sum'
) #F.binary_cross_entropy(recon_x.view(-1), x.view(-1, ), reduction='sum')
vq_loss = F.mse_loss(
z.view(-1), z_unquantized.detach().view(-1, ),
reduction='sum') + beta * F.mse_loss(
z.detach().view(-1), z_unquantized.view(-1, ), reduction='sum')
#print(KLD)
return MSE + vq_loss
func_tensor = np.vectorize(torch.from_numpy)
train_ratio = int(args.train_ratio * len(train_mora_index_lists)) #1
X_acoustic_train = [
X['acoustic']['train'][i] for i in range(len(X['acoustic']['train']))
][:train_ratio]
Y_acoustic_train = [
Y['acoustic']['train'][i] for i in range(len(Y['acoustic']['train']))
][:train_ratio]
train_mora_index_lists = [
train_mora_index_lists[i] for i in range(len(train_mora_index_lists))
][:train_ratio]
train_num = len(X_acoustic_train)
X_acoustic_test = [
X['acoustic']['test'][i] for i in range(len(X['acoustic']['test']))
]
Y_acoustic_test = [
Y['acoustic']['test'][i] for i in range(len(Y['acoustic']['test']))
]
test_mora_index_lists = [
test_mora_index_lists[i] for i in range(len(test_mora_index_lists))
]
train_loader = [[
X_acoustic_train[i], Y_acoustic_train[i], train_mora_index_lists[i]
] for i in range(len(train_mora_index_lists))]
test_loader = [[
X_acoustic_test[i], Y_acoustic_test[i], test_mora_index_lists[i]
] for i in range(len(test_mora_index_lists))]
def train(epoch):
model.train()
train_loss = 0
for batch_idx, data in enumerate(train_loader):
tmp = []
for j in range(2):
tmp.append(torch.from_numpy(data[j]).to(device))
optimizer.zero_grad()
recon_batch, z, z_unquantized = model(tmp[0], tmp[1], data[2])
loss = loss_function(recon_batch, tmp[1], z, z_unquantized)
loss.backward()
train_loss += loss.item()
optimizer.step()
del tmp
if batch_idx % len(train_loader) == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx, train_num, 100. * batch_idx / train_num,
loss.item()))
print('====> Epoch: {} Average loss: {:.4f}'.format(
epoch, train_loss / len(train_loader)))
return train_loss / len(train_loader)
def test(epoch):
model.eval()
test_loss = 0
f0_loss = 0
with torch.no_grad():
for i, data, in enumerate(test_loader):
tmp = []
for j in range(2):
tmp.append(torch.tensor(data[j]).to(device))
recon_batch, z, z_unquantized = model(tmp[0], tmp[1], data[2])
test_loss += loss_function(recon_batch, tmp[1], z,
z_unquantized).item()
f0_loss += calc_lf0_rmse(
recon_batch.cpu().numpy().reshape(-1, 199),
tmp[1].cpu().numpy().reshape(-1, 199), lf0_start_idx,
vuv_start_idx)
del tmp
test_loss /= len(test_loader)
print('====> Test set loss: {:.4f}'.format(test_loss))
return test_loss, f0_loss
loss_list = []
test_loss_list = []
test_f0_erros = []
num_epochs = args.num_epoch
for epoch in range(1, num_epochs + 1):
loss = train(epoch)
test_loss, f0_loss = test(epoch)
print('epoch [{}/{}], loss: {:.4f} test_loss: {:.4f}'.format(
epoch + 1, num_epochs, loss, test_loss))
# logging
loss_list.append(loss)
test_loss_list.append(test_loss)
test_f0_erros.append(f0_loss)
print(time.time() - start)
if epoch % 5 == 0:
torch.save(
model.state_dict(), '{}/{}layers_zdim{}_model_{}.pth'.format(
args.output_dir, num_lstm_layers, z_dim, epoch))
np.save(
args.output_dir +
'/{}layers_zdim{}_loss_list.npy'.format(num_lstm_layers, z_dim),
np.array(loss_list))
np.save(
args.output_dir + '/{}layers_zdim{}_test_loss_list.npy'.format(
num_lstm_layers, z_dim), np.array(test_loss_list))
np.save(
args.output_dir + '/{}layers_zdim{}_test_f0_loss_list.npy'.format(
num_lstm_layers, z_dim), np.array(test_f0_erros))
return f0_loss