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