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"))