def train(model, train_loader): """ Trains the C-GAE for one epoch :param epoch: the current training epoch """ model.train() losses = [] diffs = [] for batch_idx, (x, y, transform, _, _) in enumerate(train_loader): x = cuda_variable(x) y = cuda_variable(y) # Full model optimizer.zero_grad() amps_x, phases_x = model(x) amps_y, phases_y = model(y) recon_y = model.backward(amps_x, phases_y) recon_x = model.backward(amps_y, phases_x) normloss = cuda_variable(torch.Tensor(np.array([0]))) eqnloss = cuda_variable(torch.Tensor(np.array([0]))) if args.norm_loss > 0: normloss = norm_loss(model.layer.weight) * args.norm_loss if args.equal_norm_loss > 0: eqnloss = equal_norm_loss(model.layer.weight) * \ args.equal_norm_loss pow = args.power_loss loss_mse = ((y - recon_y).abs() ** pow).mean() + \ ((x - recon_x).abs() ** pow).mean() loss = loss_mse + normloss + eqnloss loss.backward() optimizer.step() if args.set_to_norm >= 0: model.set_to_norm(args.set_to_norm) losses.append(loss.item()) if args.set_to_norm >= 0: print(f"Learned norm for all bases = {model.norm_val.item()}") return losses, diffs
def run_experiment(model, start_epoch): # Train and plot intermediate results loss_curve = [] eval_curve = [] bs = args.batch_size print("Starting training...") for epoch in tqdm(range(start_epoch, args.epochs + 1)): losses, precs, eval_loss = None, None, None losses, diffs = train(model, train_loader) model.eval() eval_loss, eval_diff = eval(eval_loader) loss_curve.append(np.mean(losses)) eval_curve.append(np.mean(eval_loss)) if epoch % 1 == 0: print('Finished Epoch: {}/{} ({:.0f}%)\tLoss: {:.6f}' '\tEval: {:.6f}\tBatch Size: {}'.format( epoch, args.epochs + 1, 100. * epoch / args.epochs + 1, np.mean(losses), np.mean(eval_loss), bs)) if epoch % args.plot_interval == 0: for batch_idx, (x, y, transform, _, _) in enumerate(train_loader): if batch_idx == 0: x = cuda_variable(x) y = cuda_variable(y) plot_train_state(loss_curve, eval_curve, model, x, y, epoch) if args.data_type == 'cqt': loader = test_loader if args.test_cqt else train_loader plot_pca(loader, fn=f"pca_trans_{epoch}.png") else: break torch.save(model.state_dict(), model_save_fn) # Save the model torch.save(model.state_dict(), model_save_fn) return model
def eval(eval_loader): model.eval() losses = [] diffs = [] for batch_idx, (x, y, transform, _, _) in enumerate(eval_loader): x = cuda_variable(x) y = cuda_variable(y) amps_x, phases_x = model(x) amps_y, phases_y = model(y) recon_y = model.backward(amps_x, phases_y) recon_x = model.backward(amps_y, phases_x) pow = 1 loss_mse = (torch.abs(y - recon_y) ** pow).mean() + \ (torch.abs(x - recon_x) ** pow).mean() losses.append(loss_mse.item()) # print(losses) return losses, diffs
def to_amp_phase(input, step_size): if args.cuda: model.cuda() model.eval() ngrams = [] for i in range(0, len(input) - args.length_ngram, step_size): curr_ngram = input[i:i + args.length_ngram].reshape((-1, )) curr_ngram = standardize(curr_ngram) ngrams.append(curr_ngram) x = cuda_variable(torch.FloatTensor(np.vstack(ngrams))) ampl, phase = model(x) return ampl, phase
def plot_train_state_2d(loss_curve_eval, loss_curve_train, model, x, y, epoch, out_dir, length_ngram): model.eval() x = cuda_variable(x) y = cuda_variable(y) # calculate mapping of untransposed data amp_x, phase_x = model(x) amp_y, phase_y = model(y) recon_y = model.backward(amp_x, phase_y) recon_x = model.backward(amp_y, phase_x) make_tiles(to_numpy(recon_y).reshape(recon_y.shape[0], 1, -1, length_ngram), os.path.join(out_dir, f"recon_y{epoch}.png")) make_tiles(to_numpy(recon_x).reshape(recon_x.shape[0], 1, -1, length_ngram), os.path.join(out_dir, f"recon_x{epoch}.png")) plot_hist(to_numpy(recon_x), f"recon_x_hist_ep{epoch}", os.path.join(out_dir, f"recon_x_hist_ep{epoch}.png")) plot_hist(to_numpy(recon_y), f"recon_y_hist_ep{epoch}", os.path.join(out_dir, f"recon_y_hist_ep{epoch}.png")) half_weight = model.layer.weight.shape[0] // 2 make_tiles(to_numpy(model.layer.weight[:half_weight]).reshape(len( model.layer.weight[:half_weight]), 1, -1, length_ngram), os.path.join(out_dir, f"filters_x{epoch}.png")) make_tiles(to_numpy(model.layer.weight[half_weight:]).reshape(len( model.layer.weight[half_weight:]), 1, -1, length_ngram), os.path.join(out_dir, f"filters_y{epoch}.png")) plot_hist(to_numpy(model.layer.weight[:half_weight]).reshape(len( model.layer.weight[:half_weight]), 1, -1, length_ngram), f"filters_x_hist_ep{epoch}", os.path.join(out_dir, f"filters_x_hist_ep{epoch}.png")) plot_hist(to_numpy(model.layer.weight[half_weight:]).reshape(len( model.layer.weight[half_weight:]), 1, -1, length_ngram), f"filters_y_hist_ep{epoch}", os.path.join(out_dir, f"filters_y_hist_ep{epoch}.png")) # make_tiles(to_numpy(model.layer.weight).reshape(len( # model.layer.weight), # 1, -1, # length_ngram), # os.path.join(out_dir, f"filters_x{epoch}.png")) make_tiles(to_numpy(amp_x)[:, None, None, :], os.path.join(out_dir, f"ampx_{epoch}.png")) plot_hist(to_numpy(amp_x), f"ampx_hist_ep{epoch}", os.path.join(out_dir, f"ampx_hist_ep{epoch}.png")) make_tiles(to_numpy(amp_y)[:, None, None, :], os.path.join(out_dir, f"ampy_{epoch}.png")) plot_hist(to_numpy(amp_y), f"ampy_hist_ep{epoch}", os.path.join(out_dir, f"ampy_hist_ep{epoch}.png")) make_tiles(to_numpy(phase_x)[:, None, None, :], os.path.join(out_dir, f"phasex_{epoch}.png")) plot_hist(to_numpy(phase_x), f"phasex_hist_ep{epoch}", os.path.join(out_dir, f"phasex_hist_ep{epoch}.png")) make_tiles(to_numpy(phase_y)[:, None, None, :], os.path.join(out_dir, f"phasey_{epoch}.png")) plot_hist(to_numpy(phase_y), f"phasey_hist_ep{epoch}", os.path.join(out_dir, f"phasey_hist_ep{epoch}.png")) plot_hist(to_numpy(x.data), f"input_hist_ep{epoch}", os.path.join(out_dir, f"input_hist_ep{epoch}.png")) plot_hist(to_numpy(y.data), f"target_hist_ep{epoch}", os.path.join(out_dir, f"target_hist_ep{epoch}.png")) input_np = to_numpy(x.data) make_tiles(input_np.reshape(x.shape[0], 1, -1, length_ngram), os.path.join(out_dir, f"input_{epoch}.png")) target_np = to_numpy(y.data) make_tiles(target_np.reshape(y.shape[0], 1, -1, length_ngram), os.path.join(out_dir, f"target_{epoch}.png")) plt.clf() plt.plot(loss_curve_train) plt.plot(loss_curve_eval) plt.savefig( os.path.join(out_dir, f"loss_curve_{epoch}.png"))
def plot_train_state_1d(loss_curve_eval, loss_curve_train, model, x, y, epoch, out_dir, length_ngram): model.eval() x = cuda_variable(x) y = cuda_variable(y) # calculate mapping of untransposed data amp_x, phase_x = model(x) amp_y, phase_y = model(y) recon_y = model.backward(amp_x, phase_y) recon_x = model.backward(amp_y, phase_x) plot_audiobatch(recon_y[:20, None, :].detach().cpu(), os.path.join(out_dir, f"recon_y{epoch}.png")) plot_audiobatch(recon_x[:20, None, :].detach().cpu(), os.path.join(out_dir, f"recon_x{epoch}.png")) plot_audiobatch(y[:20, None, :].detach().cpu(), os.path.join(out_dir, f"input_y_sig{epoch}.png")) plot_audiobatch(x[:20, None, :].detach().cpu(), os.path.join(out_dir, f"input_x_sig{epoch}.png")) plot_hist(to_numpy(recon_x), f"recon_x_hist_ep{epoch}", os.path.join(out_dir, f"recon_x_hist_ep{epoch}.png")) plot_hist(to_numpy(recon_y), f"recon_y_hist_ep{epoch}", os.path.join(out_dir, f"recon_y_hist_ep{epoch}.png")) make_tiles(to_numpy(model.layer.weight).reshape(len( model.layer.weight), 1, -1, length_ngram), os.path.join(out_dir, f"filters_x{epoch}.png")) make_tiles(to_numpy(model.layer.weight).reshape(len( model.layer.weight), 1, -1, length_ngram), os.path.join(out_dir, f"filters_y{epoch}.png")) plot_audiobatch(model.layer.weight[:20, None, :].detach().cpu(), os.path.join(out_dir, f"filters_sig_real{epoch}.png")) half = model.layer.weight.shape[0] // 2 plot_audiobatch(model.layer.weight[half:half+20, None, :].detach().cpu(), os.path.join(out_dir, f"filters_sig_compl{epoch}.png")) make_tiles(to_numpy(amp_x)[:, None, None, :], os.path.join(out_dir, f"ampx_{epoch}.png")) plot_hist(to_numpy(amp_x), f"ampx_hist_ep{epoch}", os.path.join(out_dir, f"ampx_hist_ep{epoch}.png")) make_tiles(to_numpy(amp_y)[:, None, None, :], os.path.join(out_dir, f"ampy_{epoch}.png")) plot_hist(to_numpy(amp_y), f"ampy_hist_ep{epoch}", os.path.join(out_dir, f"ampy_hist_ep{epoch}.png")) make_tiles(to_numpy(phase_x)[:, None, None, :], os.path.join(out_dir, f"phasex_{epoch}.png")) plot_hist(to_numpy(phase_x), f"phasex_hist_ep{epoch}", os.path.join(out_dir, f"phasex_hist_ep{epoch}.png")) make_tiles(to_numpy(phase_y)[:, None, None, :], os.path.join(out_dir, f"phasey_{epoch}.png")) plot_hist(to_numpy(phase_y), f"phasey_hist_ep{epoch}", os.path.join(out_dir, f"phasey_hist_ep{epoch}.png")) plot_hist(to_numpy(x.data), f"input_hist_ep{epoch}", os.path.join(out_dir, f"input_hist_ep{epoch}.png")) plot_hist(to_numpy(y.data), f"target_hist_ep{epoch}", os.path.join(out_dir, f"target_hist_ep{epoch}.png")) input_np = to_numpy(x.data) make_tiles(input_np.reshape(x.shape[0], 1, -1, length_ngram), os.path.join(out_dir, f"input_{epoch}.png")) target_np = to_numpy(y.data) make_tiles(target_np.reshape(y.shape[0], 1, -1, length_ngram), os.path.join(out_dir, f"target_{epoch}.png")) plt.clf() plt.plot(loss_curve_train) plt.plot(loss_curve_eval) plt.savefig( os.path.join(out_dir, f"loss_curve_{epoch}.png"))