def train_acn(train_cnt, epoch_cnt, model_dict, data_dict, info, rescale_inv):
print('starting training routine')
base_filepath = info['base_filepath']
base_filename = os.path.split(info['base_filepath'])[1]
while train_cnt < info['num_examples_to_train']:
print('starting epoch %s on %s' % (epoch_cnt, info['device']))
train_loss_avg, train_example = run(train_cnt,
model_dict,
data_dict,
phase='train',
info=info)
epoch_cnt += 1
train_cnt += info['size_training_set']
if not epoch_cnt % info['save_every_epochs'] or epoch_cnt == 1:
# make a checkpoint
print('starting valid phase')
valid_loss_avg, valid_example = run(train_cnt,
model_dict,
data_dict,
phase='valid',
info=info)
for loss_key in train_loss_avg.keys():
for lphase in ['train_losses', 'valid_losses']:
if loss_key not in info[lphase].keys():
info[lphase][loss_key] = []
info['valid_losses'][loss_key].append(valid_loss_avg[loss_key])
info['train_losses'][loss_key].append(train_loss_avg[loss_key])
# store model
state_dict = {}
for key, model in model_dict.items():
state_dict[key + '_state_dict'] = model.state_dict()
info['train_cnts'].append(train_cnt)
info['epoch_cnt'] = epoch_cnt
state_dict['codes'] = model_dict['prior_model'].codes
state_dict['info'] = info
ckpt_filepath = os.path.join(
base_filepath, "%s_%010dex.pt" % (base_filename, train_cnt))
train_img_filepath = os.path.join(
base_filepath,
"%s_%010d_train_rec.png" % (base_filename, train_cnt))
valid_img_filepath = os.path.join(
base_filepath,
"%s_%010d_valid_rec.png" % (base_filename, train_cnt))
plot_filepath = os.path.join(
base_filepath,
"%s_%010d_loss.png" % (base_filename, train_cnt))
plot_example(train_img_filepath, train_example, num_plot=10)
plot_example(valid_img_filepath, valid_example, num_plot=10)
save_checkpoint(state_dict, filename=ckpt_filepath)
plot_losses(info['train_cnts'],
info['train_losses'],
info['valid_losses'],
name=plot_filepath,
rolling_length=1)
torch.cuda.empty_cache()
def run_autoencoder(optimizer):
""" Runs the autoencoder model using the specified optimizer.
Parameters
----------
optimizer : RMSProp/Adam
Optimization algorithm to be used for parameter learning
"""
optimizer = Adam(learning_rate=0.03) if optimizer == 'adam' else RMSProp(
learning_rate=0.05)
train_matrix, val_matrix = get_training_and_val_data()
model = Autoencoder(input_dim=train_matrix.shape[1])
model.print_summary()
model.compile(optimizer)
errors = model.fit(train_matrix,
train_matrix,
num_epochs=60,
val_set=(val_matrix, val_matrix),
early_stopping=True)
plot_losses(errors['training'], errors['validation'])
neuron_num = model.model.layers[0].optimizer.reference_index
learning_rates = model.model.layers[0].optimizer.learning_rates
plot_learning_rates(learning_rates['weights'], learning_rates['bias'],
neuron_num)
def fit(args, data, val_data, model):
# dtype = torch.FloatTensor or #change to torch.cuda.FloatTensor to make it run on GPU
dataloader = torch_utils.DataLoader(data, batch_size=args.batch_size, shuffle=True)
val_dataloader = torch_utils.DataLoader(val_data, batch_size=args.batch_size, shuffle=True)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
plot_loss = []
plot_val_loss = []
for epoch in range(args.epochs):
for i, batch in enumerate(dataloader):
N, L = batch.shape
batch = Variable(batch.float(), requires_grad=False)
output = model(batch.view(N, 1, L))
loss = criterion(output, target=batch)
if i % 100 == 0:
val_batch = next(iter(val_dataloader))
val_batch = Variable(val_batch.float(), requires_grad=False)
val_output = model(val_batch.view(N, 1, L))
val_loss = criterion(val_output, target=val_batch)
plot_loss.append(loss.data.numpy()[0])
plot_val_loss.append(val_loss.data.numpy()[0])
print 'epoch', epoch, 'num', i, 'loss', loss.data.numpy()[0], 'val loss', val_loss.data.numpy()[0]
optimizer.zero_grad() # zero the gradient buffers
loss.backward()
optimizer.step() # Does the update
plot_losses(plot_loss, plot_val_loss)
def train(name, gen, disc, gen_train_ratio=5, epochs=-1):
device = utils.get_device()
dataloader = data.get_dataloader()
loss_func = nn.BCELoss()
g_optimizer = to.Adam(gen.parameters(), lr=.0003, betas=(.5, .9))
d_optimizer = to.Adam(disc.parameters(), lr=.0003, betas=(.5, .9))
iter_axis = []
g_loss_axis = []
d_loss_axis = []
fixed_noise = torch.randn(64, gen.nz, device=device)
epoch_iterator = range(epochs) if epochs >= 0 else count(0)
iters = 0
for epoch in epoch_iterator:
for i, batch in enumerate(dataloader):
loss_disc, acc_real, acc_fake = train_discriminator(
gen, disc, batch, loss_func, d_optimizer)
for _ in range(gen_train_ratio):
loss_gen, acc_gen = train_generator(gen, disc, loss_func,
g_optimizer)
# Training stats
if i % 50 == 0:
print(
f'[{epoch:2d}/{epochs}][{i:3d}/{len(dataloader)}]\t' +\
f'Loss_D: {loss_disc:3.4f}\tLoss_G: {loss_gen:3.4f}\t' +\
f'D(x): {acc_real:3.4f}\tD(G(z)): {acc_fake:3.4f} / {acc_gen:3.4f}'
)
if (iters % 10 == 0) or ((epoch == epochs - 1) and
(i == len(dataloader) - 1)):
iter_axis.append(iters)
g_loss_axis.append(loss_gen)
d_loss_axis.append(loss_disc)
# Save output
if (iters % 500 == 0) or ((epoch == epochs - 1) and
(i == len(dataloader) - 1)):
print('Saving images...')
with torch.no_grad():
gen.mixing = False
fake = gen(fixed_noise).detach().cpu()
gen.mixing = True
utils.save_image(utils.make_torch_grid(fake),
f'{name}_{iters}.png')
# Save weights
if (iters % 1000 == 0) or ((epoch == epochs - 1) and
(i == len(dataloader) - 1)):
print('Saving weights...')
utils.save_weights(gen, f'{name}_gen.pt')
utils.save_weights(disc, f'{name}_disc.pt')
iters += 1
utils.plot_losses(iter_axis, g_loss_axis, d_loss_axis)
def find_batch_z(gen, x, nz, lr, exDir, maxEpochs=100, alpha=1e-6, batchNo=0):
#generator in eval mode
gen.eval()
#save the "original" images
save_image(x.data,
join(exDir, 'original_batch' + str(batchNo) + '.png'),
normalize=True)
#Assume the prior is Standard Normal
pdf = torch.distributions.Normal(0, 1)
if gen.useCUDA:
Zinit = Variable(torch.randn(x.size(0), opts.nz).cuda(),
requires_grad=True)
else:
Zinit = Variable(torch.randn(x.size(0), opts.nz), requires_grad=True)
#optimizer
optZ = torch.optim.RMSprop([Zinit], lr=lr)
losses = {'rec': [], 'logProb': []}
for e in range(maxEpochs):
#reconstruction loss
xHAT = gen.forward(Zinit)
recLoss = F.mse_loss(xHAT, x)
#loss to make sure z's are Guassian
logProb = pdf.log_prob(Zinit).mean(
dim=1
) #each element of Z is independant, so likelihood is a sum of log of elements
loss = recLoss - (alpha * logProb.mean())
optZ.zero_grad()
loss.backward()
optZ.step()
losses['rec'].append(recLoss.data[0])
losses['logProb'].append(logProb.mean().data[0])
if e % 100 == 0:
print '[%d] loss: %0.5f, recLoss: %0.5f, regMean: %0.5f' % (
e, loss.data[0], recLoss.data[0], logProb.mean().data[0])
# save_image(xHAT.data, join(exDir, 'rec'+str(e)+'.png'), normalize=True)
#plot training losses
if e > 0:
plot_losses(losses, exDir, e + 1)
plot_norm_losses(losses, exDir, e + 1)
#visualise the final output
xHAT = gen.forward(Zinit)
save_image(xHAT.data,
join(exDir, 'rec_batch' + str(batchNo) + '.png'),
normalize=True)
return Zinit, recLoss.data[0], xHAT
def train(self):
self.train_hist = dict()
self.train_hist["D_loss"] = list()
self.train_hist["G_loss"] = list()
y_real = Variable(torch.ones((self.batch_size, 1)))
y_fake = Variable(torch.zeros((self.batch_size, 1)))
for epoch in range(self.epoch):
self.G.train()
for step, (x_real, _) in enumerate(self.data_loader):
if step == self.data_loader.__len__() - 1:
# ignore the last batch
break
z = Variable(torch.rand(self.batch_size, self.inp_dim))
x_real = Variable(x_real)
# train Discriminator
self.D_optimizer.zero_grad()
# loss when D is identifying real dataset(should predict 1)
D_real = self.D(x_real)
D_real_loss = self.criterion(D_real, y_real)
# loss when D is identifying fake dataset(should predict 0)
x_fake = self.G(z)
D_fake = self.D(x_fake)
D_fake_loss = self.criterion(D_fake, y_fake)
D_loss = D_fake_loss + D_real_loss
self.train_hist["D_loss"].append(D_loss.data[0])
D_loss.backward()
self.D_optimizer.step()
# train Generator
self.G_optimizer.zero_grad()
x_generated = self.G(z)
D_fake = self.D(x_generated)
# Generator should aim to make prediction of Discriminator
# close to 1
G_loss = self.criterion(D_fake, y_real)
self.train_hist["G_loss"].append(G_loss.data[0])
G_loss.backward()
self.G_optimizer.step()
if (step + 1) % 100 == 0:
print("Epoch: [%2d] D_loss: %.5f G_loss: %.5f" %
(epoch + 1, D_loss.data[0], G_loss.data[0]))
self.visualize(epoch + 1)
print("Training complete!!")
utils.plot_losses(self.train_hist)
def loop(data_loader, num_epochs=1000, save_every=1000, train_losses=[], test_losses=[], train_cnts=[], test_cnts=[], dummy=False):
print("starting training loop for data with %s batches"%data_loader.num_batches)
st = time.time()
if len(train_losses):
# resume cnt from last save
last_save = train_cnts[-1]
cnt = train_cnts[-1]
else:
last_save = 0
cnt = 0
v_xn, v_yn = data_loader.validation_data()
v_x = Variable(torch.FloatTensor(np.swapaxes(v_xn,1,0))).to(DEVICE)
v_y = Variable(torch.FloatTensor(np.swapaxes(v_yn,1,0))).to(DEVICE)
if dummy:
print("WARNING DUMMMY Validation")
v_x, v_y = get_dummy_data(v_x, v_y)
for e in range(num_epochs):
ecnt = 0
tst = round((time.time()-st)/60., 0)
if not e%1 and e>0:
print("starting epoch %s, %s mins, loss %s, seen %s, last save at %s" %(e, tst, train_losses[-1], cnt, last_save))
batch_loss = []
for b in range(data_loader.num_batches):
x, y = data_loader.next_batch()
x = Variable(torch.FloatTensor(np.swapaxes(x,1,0))).to(DEVICE)
y = Variable(torch.FloatTensor(np.swapaxes(y,1,0))).to(DEVICE)
if dummy:
y_pred, loss = train(v_x, v_y, validation=False)
print('DUMMY test loss', cnt, loss)
else:
y_pred, loss = train(x.to(DEVICE),y.to(DEVICE),validation=False)
train_cnts.append(cnt)
train_losses.append(loss)
if cnt%100:
valy_pred, val_mean_loss = train(v_x,v_y,validation=True)
test_losses.append(val_mean_loss)
test_cnts.append(cnt)
if cnt-last_save >= save_every:
last_save = cnt
# find test loss
print('epoch: {} saving after example {} train loss {} test loss {}'.format(e,cnt,loss,val_mean_loss))
state = {
'train_cnts':train_cnts,
'train_losses':train_losses,
'test_cnts': test_cnts,
'test_losses':test_losses,
'state_dict':lstm.state_dict(),
'optimizer':optim.state_dict(),
}
basename = os.path.join(savedir, '%s_%015d'%(model_save_name,cnt))
plot_losses(train_cnts, train_losses, test_cnts, test_losses, name=basename+'_loss.png')
save_checkpoint(state, filename=basename+'.pkl')
cnt+= x.shape[1]
ecnt+= x.shape[1]
def train(train_data, video_train_data, eval_data, video_eval_data, test_data, video_test_data, model, batch_size, num_epochs, model_name):
model = model.train()
weights = [680/261, 680/419]
class_weights = torch.tensor(weights, device=device)
criterion = nn.CrossEntropyLoss(weight=class_weights).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, factor=0.5, cooldown=0, min_lr=1e-6, verbose=True)
best_valid_loss = float('inf')
best_eval_acc = 0.0
eval_pat = 10
dev_loss_list = []
train_loss_list = []
train_acc_list = []
dev_acc_list = []
step_num = 0
for epoch_id in range(num_epochs):
for batch_features, video_batch_features, batch_labels in read_batches(train_data, video_train_data, batch_size, True, device, is_multi=False):
model = model.train()
optimizer.zero_grad()
step_num+=1
batch_features = batch_features.permute((0,2,1))
out = model(video_batch_features)
outs = [out]
losses = []
for i, out in enumerate(outs):
losses.append(criterion(
out, # (batch_size , num_classes)
batch_labels[:, i:i + 1].view(-1) # (batch_size * 1)
))
loss = sum(losses)
loss.backward()
optimizer.step()
if step_num % eval_pat == 0:
# print("plot now")
train_loss, train_bal, train_f, train_acc = evaluate(model, train_data, video_train_data, batch_size, criterion)
dev_loss, dev_bal, dev_f, dev_acc = evaluate(model, eval_data, video_eval_data, batch_size, criterion)
train_loss_list.append(train_loss)
dev_loss_list.append(dev_loss)
train_acc_list.append(train_bal)
dev_acc_list.append(dev_bal)
plot_losses(model_name, train_loss_list, dev_loss_list, train_acc_list, dev_acc_list, step_num)
train_loss, train_bal, train_f, train_acc = evaluate(model, train_data, video_train_data, batch_size, criterion)
dev_loss, dev_bal, dev_f, dev_acc = evaluate(model, eval_data, video_eval_data, batch_size, criterion)
scheduler.step(dev_loss)
print()
print(f'TRAIN Epoch {epoch_id} | balanced train. accuracy={train_bal} | train fscore={train_f} | train_loss={ff(train_loss,n=8)} | train. accuracy={train_acc} |')
print(f'Epoch {epoch_id} | balanced dev. accuracy={dev_bal} | dev fscore={dev_f} | dev_loss={ff(dev_loss, n=8)} | dev. accuracy={dev_acc} |')
if dev_loss < best_valid_loss:
best_valid_loss = dev_loss
torch.save(model.state_dict(), model_name+'.pt')
model.load_state_dict(torch.load(model_name+'.pt'))
dev_loss, dev_bal, dev_f, dev_acc = evaluate(model, test_data, video_test_data, batch_size, criterion)
print(f'TEST - balanced test. accuracy={dev_bal} | test fscore={dev_f} | test_loss={ff(dev_loss, n=8)} | test acc={dev_acc} |')
plot_losses(model_name, train_loss_list, dev_loss_list, train_acc_list, dev_acc_list, step_num, test_acc=dev_bal, test_loss=dev_loss)
def main(args):
"""main procedure"""
# get configuration
device = global_conf["device"]
image_size = global_conf["image_size"]
data_dir = global_conf["data_dir"]
res_dir = global_conf["res_dir"]
save_dir = global_conf["checkpoints_dir"]
# prepare data
train_loader, test_loader = prepare_data(args, dir_path=data_dir)
# prepare model
model = ConvJointVAE(args.temperature, args.disc_dims, args.disc_cap,
args.cont_cap, image_size, args.channels,
args.hidden_size, args.dim_z)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
print(model)
# train and test
losses = {}
for epoch in range(1, args.epochs + 1):
avg_loss = train(model, train_loader, epoch, optimizer, args, device)
# record all the losses
for k in avg_loss.keys():
if k in losses:
losses[k].append(avg_loss[k])
else:
losses[k] = [avg_loss[k]]
test(model, test_loader, epoch, args, device, res_dir)
with torch.no_grad():
sample = model.sample(64, device).cpu()
save_image(sample, res_dir + '/sample_' + str(epoch) + '.png')
# plot train losses
plot_losses(losses, res_dir + '/loss.png')
# save the model and related params
if args.save:
save_dir = os.path.join(save_dir, 'mnist')
save_(model,
save_dir,
args,
global_conf,
comment=args.tag,
extra={
"cont_cap": model.cont_cap_current,
"disc_cap": model.disc_cap_current
})
def train_iters(model, optimizer, epochs, current_epoch, batch_size,
print_every_batch, save_every_epoch, save_dir):
all_losses = []
for epoch in range(current_epoch, epochs + 1):
model.train()
num_batches = (len(train_sentences) - 1) // batch_size + 1
current_loss = 0
for batch_i in range(num_batches):
optimizer.zero_grad()
start_idx = batch_i * batch_size
end_idx = (batch_i + 1) * batch_size
sentences, tags = train_sentences[start_idx:end_idx], train_tags[
start_idx:end_idx]
loss = -model(sentences, tags)
loss.backward()
optimizer.step()
loss_value = loss.item() / sum(
[len(sentence) for sentence in sentences])
# loss_value = loss.item()
all_losses.append(loss_value)
current_loss += loss_value
if (batch_i + 1) % print_every_batch == 0:
print(
f"Epoch {epoch}, batch {batch_i}: loss = {current_loss / print_every_batch}"
)
current_loss = 0
if epoch % save_every_epoch == 0:
directory = os.path.join(
save_dir, f'{character_hidden_dim}_{context_hidden_dim}')
if not os.path.exists(directory):
os.makedirs(directory)
torch.save(
{
'epoch': epoch,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(directory, f"{epoch}_checkpoint.tar"))
print('Evaluating...')
eval_result = eval2(model)
print(f"Epoch {epoch}: {eval_result}")
utils.plot_losses(all_losses)
def run(args):
# Images
style_img, content_img, gen_img = utils.get_starting_imgs(args)
# CNN layers
style_layers, content_layers = cnn.get_layers(args)
# Make model
model = arch.make_model(args, style_layers, content_layers, style_img,
content_img)
# Transfer
losses_dict, gen_hist = style.transfer(args, gen_img, style_img, model)
# Plot losses
loss_fig = utils.plot_losses(losses_dict)
# Save
# Resized style
utils.save_tensor_img(style_img, os.path.join(args.out_dir, 'style.png'))
# Generated image
utils.save_tensor_img(gen_img, os.path.join(args.out_dir, 'gen.png'))
gen_hist[0].save(os.path.join(args.out_dir, 'gen.gif'),
save_all=True,
append_images=gen_hist[1:])
# Losses
loss_fig.savefig(os.path.join(args.out_dir, 'losses.png'))
print(f"Results saved to '{args.out_dir}'")
def main():
df, X, y = preprocess_data()
X_train, X_test, y_train, y_test = train_test_splitter(X=X, y=y, ratio=0.8)
logistic_regressor = LogisticRegressor(alpha=0.05,
c=0.01,
T=1000,
random_seed=0,
intercept=True)
losses = logistic_regressor.fit(X_train, y_train)
plot_losses(losses=losses, savefig=True)
train_error = error_rate(y_train, logistic_regressor.predict(X_train))
test_error = error_rate(y_test, logistic_regressor.predict(X_test))
print('Training Error Rate: %f' % train_error)
print('Test Error Rate: %f' % test_error)
def train_loop(self, num_batches=10):
ecnt = 0
batch_loss = []
for b in range(data_loader.num_batches):
xnp, ynp = self.data_loader.next_batch()
x = Variable(torch.FloatTensor(xnp))
y = Variable(torch.FloatTensor(ynp))
y_pred, loss = train(x,y,validation=False)
train_cnts.append(cnt)
train_losses.append(loss)
if cnt%100:
valy_pred, val_mean_loss = train(v_x,v_y,validation=True)
test_losses.append(val_mean_loss)
test_cnts.append(cnt)
if cnt-last_save >= save_every:
last_save = cnt
# find test loss
print('epoch: {} saving after example {} train loss {} test loss {}'.format(e,cnt,loss,val_mean_loss))
state = {
'train_cnts':train_cnts,
'train_losses':train_losses,
'test_cnts': test_cnts,
'test_losses':test_losses,
'state_dict':lstm.state_dict(),
'optimizer':optim.state_dict(),
}
basename = os.path.join(savedir, '%s_%015d'%(model_save_name,cnt))
n = 500
plot_losses(rolling_average(train_cnts, n),
rolling_average(train_losses, n),
rolling_average(test_cnts, n),
rolling_average(test_losses, n), name=basename+'_loss.png')
save_checkpoint(state, filename=basename+'.pkl')
cnt+= x.shape[1]
ecnt+= x.shape[1]
loop(data_loader, save_every=save_every, num_epochs=args.num_epochs,
train_losses=train_losses, test_losses=test_losses,
train_cnts=train_cnts, test_cnts=test_cnts, dummy=args.dummy)
def find_z(gen, x, nz, lr, exDir, maxEpochs=100):
#generator in eval mode
gen.eval()
#save the "original" images
save_image(x.data, join(exDir, 'original.png'), normalize=True)
if gen.useCUDA:
gen.cuda()
Zinit = Variable(torch.randn(x.size(0), opts.nz).cuda(),
requires_grad=True)
else:
Zinit = Variable(torch.randn(x.size(0), opts.nz), requires_grad=True)
#optimizer
optZ = torch.optim.RMSprop([Zinit], lr=lr)
losses = {'rec': []}
for e in range(maxEpochs):
xHAT = gen.forward(Zinit)
recLoss = F.mse_loss(xHAT, x)
optZ.zero_grad()
recLoss.backward()
optZ.step()
losses['rec'].append(recLoss.data[0])
print '[%d] loss: %0.5f' % (e, recLoss.data[0])
#plot training losses
if e > 0:
plot_losses(losses, exDir, e + 1)
#visualise the final output
xHAT = gen.forward(Zinit)
save_image(xHAT.data, join(exDir, 'rec.png'))
return Zinit
def main(args):
"""main procedure"""
# get configuration
device = global_conf["device"]
img_size = global_conf["image_size"]
data_dir = global_conf["data_dir"]
res_dir = global_conf["res_dir"]
save_dir = global_conf["checkpoints_dir"]
# prepare data
train_loaders, test_loaders = prepare_data(args, dir_path=data_dir)
# prepare model
model = FactorVAE(img_size[0]*img_size[1], args.n_hidden, args.dim_z,
img_size[0]*img_size[1], args.gamma)
optimizer_vae = optim.Adam(model.parameters(), lr=args.lr)
optimizer_discriminator = optim.Adam(model.discriminator.parameters(), lr=args.lr)
optimizers = {"vae": optimizer_vae, "discriminator": optimizer_discriminator}
# train and test
losses_vae = []
losses_discriminator = []
for epoch in range(1, args.epochs+1):
avg_loss = train(model, train_loaders, epoch,
optimizers, args, device, img_size)
losses_vae.append(avg_loss["vae"])
losses_discriminator.append(avg_loss["discriminator"])
test(model, test_loaders, epoch, args, device, img_size, res_dir)
with torch.no_grad():
sample = model.sample(64, device).cpu()
save_image(sample.view(
64, 1, img_size[0], img_size[1]), res_dir+'/sample_'+str(epoch)+'.png')
# plot train losses
plot_losses({"vae": losses_vae, "discriminator": losses_discriminator}, res_dir+'/loss.png')
# save the model and related params
if args.save:
save_dir = os.path.join(save_dir, 'mnist')
save_(model, save_dir, args, global_conf, comment=args.tag)
def main():
criterion = torch.nn.MSELoss()
noise_values = torch.linspace(0.1, 2.5, 10)
noise_loss = []
for noise in noise_values:
data_maker = DataMaker(noise=noise)
X_train, y_train = data_maker.get_train_data()
X_test, y_test = data_maker.get_test_data()
model = LSTM(input_dim=INPUT_SIZE,
hidden_dim=SEQ_SIZE,
batch_size=BATCH_SIZE,
output_dim=OUTPUT_DIM,
num_layers=NUM_LAYERS)
model, training_losses, test_losses = train(X_train=X_train,
y_train=y_train,
X_test=X_test,
y_test=y_test,
model=model,
batch_size=BATCH_SIZE)
plot_losses(title='Sequence LSTM noise = {}'.format(
round(noise.item(), 3)),
train_loss=training_losses,
test_loss=test_losses,
epochs=range(NUM_EPOCHS))
noise_loss += [criterion(model(X_test), y_test)]
plt.style.use("ggplot")
plt.plot(noise_values, noise_loss, 'b', label="Test Loss")
plt.xlabel("Noise")
plt.ylabel("Loss")
plt.legend(loc="upper left")
plt.title('Sequence LSTM - noises')
plt.show()
def run(args):
# Images
style_img, content_img, gen_img = utils.get_starting_imgs(args)
# CNN layers
style_layers, content_layers = cnn.get_layers(args)
# Make model
style_model = transfer_model.make(args, style_layers, content_layers,
style_img, content_img)
# Transfer
losses_dict = style.transfer(args, gen_img, style_img, style_model)
# Losses
loss_fig = utils.plot_losses(losses_dict)
# Save generated image
utils.save_tensor_img(gen_img, os.path.join(args.out_dir, 'gen.png'))
# Save losses
loss_fig.savefig(os.path.join(args.out_dir, 'losses.pdf'))
def game6():
situation_size, message_size, prediction_size, func_size, hidden_size = (
10,
2,
10,
20,
64,
)
num_reproductions = 10
all_losses = []
for _ in range(num_reproductions):
game = Game(situation_size, message_size, prediction_size, func_size,
hidden_size, -1)
print_first = True
for lr in [0.01, 0.001, 0.0001]:
play_game(game,
1000,
learning_rate=lr,
func_out_training=[0, 1, 2, 3, 5, 7])
if print_first:
logging.info(
f"Epoch {game.loss_per_epoch[0][0]}:\t{game.loss_per_epoch[0][1]:.2e}"
)
print_first = False
logging.info(
f"Epoch {game.loss_per_epoch[-1][0]}:\t{game.loss_per_epoch[-1][1]:.2e}"
)
all_losses.append(get_loss_per_function(game))
# plot_messages_information(game, 40)
all_losses = np.array(all_losses)
loss_average_per_func = np.average(all_losses, axis=0)
A = plot_losses(game, loss_average_per_func)
from utils import prepare_data, prepare_model, plot_losses, evaluate_model
from keras.callbacks import ModelCheckpoint, TensorBoard
import random
#Creates our training/val/testing splits
random.seed(9001)
X_train, X_val, X_test, y_train, y_val, y_test, input_output = prepare_data('./recordings/')
#Prepares a CNN with our desired architecture
CNN_best_model = prepare_model(input_output, modeltype='CNN', dropout=False, maxpooling=True, batch_n=False)
#Creates callbacks to save best model in training
callbacks = [ModelCheckpoint(filepath='models/cnn_best_model.h5', monitor='val_loss', save_best_only=True), TensorBoard(log_dir='./Graph', histogram_freq=1,
write_graph=False, write_images=False)]
#Fits model
history = CNN_best_model.fit(X_train, y_train, batch_size=32, epochs=50, verbose= 2, validation_data = [X_val, y_val],
callbacks=callbacks)
#Plots loss curve
plot_losses(history)
#Evaluate model on testing set
evaluate_model('models/cnn_best_model.h5', X_test, y_test)
def train():
# Fix Seed for Reproducibility #
torch.manual_seed(9)
if torch.cuda.is_available():
torch.cuda.manual_seed(9)
# Samples, Weights and Results Path #
paths = [config.samples_path, config.weights_path, config.plots_path]
paths = [make_dirs(path) for path in paths]
# Prepare Data Loader #
train_horse_loader, train_zebra_loader = get_horse2zebra_loader('train', config.batch_size)
val_horse_loader, val_zebra_loader = get_horse2zebra_loader('test', config.batch_size)
total_batch = min(len(train_horse_loader), len(train_zebra_loader))
# Image Pool #
masked_fake_A_pool = ImageMaskPool(config.pool_size)
masked_fake_B_pool = ImageMaskPool(config.pool_size)
# Prepare Networks #
Attn_A = Attention()
Attn_B = Attention()
G_A2B = Generator()
G_B2A = Generator()
D_A = Discriminator()
D_B = Discriminator()
networks = [Attn_A, Attn_B, G_A2B, G_B2A, D_A, D_B]
for network in networks:
network.to(device)
# Loss Function #
criterion_Adversarial = nn.MSELoss()
criterion_Cycle = nn.L1Loss()
# Optimizers #
D_optim = torch.optim.Adam(chain(D_A.parameters(), D_B.parameters()), lr=config.lr, betas=(0.5, 0.999))
G_optim = torch.optim.Adam(chain(Attn_A.parameters(), Attn_B.parameters(), G_A2B.parameters(), G_B2A.parameters()), lr=config.lr, betas=(0.5, 0.999))
D_optim_scheduler = get_lr_scheduler(D_optim)
G_optim_scheduler = get_lr_scheduler(G_optim)
# Lists #
D_A_losses, D_B_losses = [], []
G_A_losses, G_B_losses = [], []
# Train #
print("Training Unsupervised Attention-Guided GAN started with total epoch of {}.".format(config.num_epochs))
for epoch in range(config.num_epochs):
for i, (real_A, real_B) in enumerate(zip(train_horse_loader, train_zebra_loader)):
# Data Preparation #
real_A = real_A.to(device)
real_B = real_B.to(device)
# Initialize Optimizers #
D_optim.zero_grad()
G_optim.zero_grad()
###################
# Train Generator #
###################
set_requires_grad([D_A, D_B], requires_grad=False)
# Adversarial Loss using real A #
attn_A = Attn_A(real_A)
fake_B = G_A2B(real_A)
masked_fake_B = fake_B * attn_A + real_A * (1-attn_A)
masked_fake_B *= attn_A
prob_real_A = D_A(masked_fake_B)
real_labels = torch.ones(prob_real_A.size()).to(device)
G_loss_A = criterion_Adversarial(prob_real_A, real_labels)
# Adversarial Loss using real B #
attn_B = Attn_B(real_B)
fake_A = G_B2A(real_B)
masked_fake_A = fake_A * attn_B + real_B * (1-attn_B)
masked_fake_A *= attn_B
prob_real_B = D_B(masked_fake_A)
real_labels = torch.ones(prob_real_B.size()).to(device)
G_loss_B = criterion_Adversarial(prob_real_B, real_labels)
# Cycle Consistency Loss using real A #
attn_ABA = Attn_B(masked_fake_B)
fake_ABA = G_B2A(masked_fake_B)
masked_fake_ABA = fake_ABA * attn_ABA + masked_fake_B * (1 - attn_ABA)
# Cycle Consistency Loss using real B #
attn_BAB = Attn_A(masked_fake_A)
fake_BAB = G_A2B(masked_fake_A)
masked_fake_BAB = fake_BAB * attn_BAB + masked_fake_A * (1 - attn_BAB)
# Cycle Consistency Loss #
G_cycle_loss_A = config.lambda_cycle * criterion_Cycle(masked_fake_ABA, real_A)
G_cycle_loss_B = config.lambda_cycle * criterion_Cycle(masked_fake_BAB, real_B)
# Total Generator Loss #
G_loss = G_loss_A + G_loss_B + G_cycle_loss_A + G_cycle_loss_B
# Back Propagation and Update #
G_loss.backward()
G_optim.step()
#######################
# Train Discriminator #
#######################
set_requires_grad([D_A, D_B], requires_grad=True)
# Train Discriminator A using real A #
prob_real_A = D_A(real_B)
real_labels = torch.ones(prob_real_A.size()).to(device)
D_loss_real_A = criterion_Adversarial(prob_real_A, real_labels)
# Add Pooling #
masked_fake_B, attn_A = masked_fake_B_pool.query(masked_fake_B, attn_A)
masked_fake_B *= attn_A
# Train Discriminator A using fake B #
prob_fake_B = D_A(masked_fake_B.detach())
fake_labels = torch.zeros(prob_fake_B.size()).to(device)
D_loss_fake_A = criterion_Adversarial(prob_fake_B, fake_labels)
D_loss_A = (D_loss_real_A + D_loss_fake_A).mean()
# Train Discriminator B using real B #
prob_real_B = D_B(real_A)
real_labels = torch.ones(prob_real_B.size()).to(device)
D_loss_real_B = criterion_Adversarial(prob_real_B, real_labels)
# Add Pooling #
masked_fake_A, attn_B = masked_fake_A_pool.query(masked_fake_A, attn_B)
masked_fake_A *= attn_B
# Train Discriminator B using fake A #
prob_fake_A = D_B(masked_fake_A.detach())
fake_labels = torch.zeros(prob_fake_A.size()).to(device)
D_loss_fake_B = criterion_Adversarial(prob_fake_A, fake_labels)
D_loss_B = (D_loss_real_B + D_loss_fake_B).mean()
# Calculate Total Discriminator Loss #
D_loss = D_loss_A + D_loss_B
# Back Propagation and Update #
D_loss.backward()
D_optim.step()
# Add items to Lists #
D_A_losses.append(D_loss_A.item())
D_B_losses.append(D_loss_B.item())
G_A_losses.append(G_loss_A.item())
G_B_losses.append(G_loss_B.item())
####################
# Print Statistics #
####################
if (i+1) % config.print_every == 0:
print("UAG-GAN | Epoch [{}/{}] | Iteration [{}/{}] | D A Losses {:.4f} | D B Losses {:.4f} | G A Losses {:.4f} | G B Losses {:.4f}".
format(epoch+1, config.num_epochs, i+1, total_batch, np.average(D_A_losses), np.average(D_B_losses), np.average(G_A_losses), np.average(G_B_losses)))
# Save Sample Images #
save_samples(val_horse_loader, val_zebra_loader, G_A2B, G_B2A, Attn_A, Attn_B, epoch, config.samples_path)
# Adjust Learning Rate #
D_optim_scheduler.step()
G_optim_scheduler.step()
# Save Model Weights #
if (epoch + 1) % config.save_every == 0:
torch.save(G_A2B.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Generator_A2B_Epoch_{}.pkl'.format(epoch+1)))
torch.save(G_B2A.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Generator_B2A_Epoch_{}.pkl'.format(epoch+1)))
torch.save(Attn_A.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Attention_A_Epoch_{}.pkl'.format(epoch+1)))
torch.save(Attn_B.state_dict(), os.path.join(config.weights_path, 'UAG-GAN_Attention_B_Epoch_{}.pkl'.format(epoch+1)))
# Make a GIF file #
make_gifs_train("UAG-GAN", config.samples_path)
# Plot Losses #
plot_losses(D_A_losses, D_B_losses, G_A_losses, G_B_losses, config.num_epochs, config.plots_path)
print("Training finished.")
def train():
# Fix Seed for Reproducibility #
torch.manual_seed(9)
if torch.cuda.is_available():
torch.cuda.manual_seed(9)
# Samples, Weights, and Plots Path #
paths = [config.samples_path, config.weights_path, config.plots_path]
paths = [make_dirs(path) for path in paths]
# Prepare Data Loader #
train_loader_selfie, train_loader_anime = get_selfie2anime_loader(
'train', config.batch_size)
total_batch = max(len(train_loader_selfie), len(train_loader_anime))
test_loader_selfie, test_loader_anime = get_selfie2anime_loader(
'test', config.val_batch_size)
# Prepare Networks #
D_A = Discriminator(num_layers=7)
D_B = Discriminator(num_layers=7)
L_A = Discriminator(num_layers=5)
L_B = Discriminator(num_layers=5)
G_A2B = Generator(image_size=config.crop_size,
num_blocks=config.num_blocks)
G_B2A = Generator(image_size=config.crop_size,
num_blocks=config.num_blocks)
networks = [D_A, D_B, L_A, L_B, G_A2B, G_B2A]
for network in networks:
network.to(device)
# Loss Function #
Adversarial_loss = nn.MSELoss()
Cycle_loss = nn.L1Loss()
BCE_loss = nn.BCEWithLogitsLoss()
# Optimizers #
D_optim = torch.optim.Adam(chain(D_A.parameters(), D_B.parameters(),
L_A.parameters(), L_B.parameters()),
lr=config.lr,
betas=(0.5, 0.999),
weight_decay=0.0001)
G_optim = torch.optim.Adam(chain(G_A2B.parameters(), G_B2A.parameters()),
lr=config.lr,
betas=(0.5, 0.999),
weight_decay=0.0001)
D_optim_scheduler = get_lr_scheduler(D_optim)
G_optim_scheduler = get_lr_scheduler(G_optim)
# Rho Clipper to constraint the value of rho in AdaILN and ILN #
Rho_Clipper = RhoClipper(0, 1)
# Lists #
D_losses = []
G_losses = []
# Train #
print("Training U-GAT-IT started with total epoch of {}.".format(
config.num_epochs))
for epoch in range(config.num_epochs):
for i, (selfie, anime) in enumerate(
zip(train_loader_selfie, train_loader_anime)):
# Data Preparation #
real_A = selfie.to(device)
real_B = anime.to(device)
# Initialize Optimizers #
D_optim.zero_grad()
G_optim.zero_grad()
#######################
# Train Discriminator #
#######################
set_requires_grad([D_A, D_B, L_A, L_B], requires_grad=True)
# Forward Data #
fake_B, _, _ = G_A2B(real_A)
fake_A, _, _ = G_B2A(real_B)
G_real_A, G_real_A_cam, _ = D_A(real_A)
L_real_A, L_real_A_cam, _ = L_A(real_A)
G_real_B, G_real_B_cam, _ = D_B(real_B)
L_real_B, L_real_B_cam, _ = L_B(real_B)
G_fake_A, G_fake_A_cam, _ = D_A(fake_A)
L_fake_A, L_fake_A_cam, _ = L_A(fake_A)
G_fake_B, G_fake_B_cam, _ = D_B(fake_B)
L_fake_B, L_fake_B_cam, _ = L_B(fake_B)
# Adversarial Loss of Discriminator #
real_labels = torch.ones(G_real_A.shape).to(device)
D_ad_real_loss_GA = Adversarial_loss(G_real_A, real_labels)
fake_labels = torch.zeros(G_fake_A.shape).to(device)
D_ad_fake_loss_GA = Adversarial_loss(G_fake_A, fake_labels)
D_ad_loss_GA = D_ad_real_loss_GA + D_ad_fake_loss_GA
real_labels = torch.ones(G_real_A_cam.shape).to(device)
D_ad_cam_real_loss_GA = Adversarial_loss(G_real_A_cam, real_labels)
fake_labels = torch.zeros(G_fake_A_cam.shape).to(device)
D_ad_cam_fake_loss_GA = Adversarial_loss(G_fake_A_cam, fake_labels)
D_ad_cam_loss_GA = D_ad_cam_real_loss_GA + D_ad_cam_fake_loss_GA
real_labels = torch.ones(G_real_B.shape).to(device)
D_ad_real_loss_GB = Adversarial_loss(G_real_B, real_labels)
fake_labels = torch.zeros(G_fake_B.shape).to(device)
D_ad_fake_loss_GB = Adversarial_loss(G_fake_B, fake_labels)
D_ad_loss_GB = D_ad_real_loss_GB + D_ad_fake_loss_GB
real_labels = torch.ones(G_real_B_cam.shape).to(device)
D_ad_cam_real_loss_GB = Adversarial_loss(G_real_B_cam, real_labels)
fake_labels = torch.zeros(G_fake_B_cam.shape).to(device)
D_ad_cam_fake_loss_GB = Adversarial_loss(G_fake_B_cam, fake_labels)
D_ad_cam_loss_GB = D_ad_cam_real_loss_GB + D_ad_cam_fake_loss_GB
# Adversarial Loss of L #
real_labels = torch.ones(L_real_A.shape).to(device)
D_ad_real_loss_LA = Adversarial_loss(L_real_A, real_labels)
fake_labels = torch.zeros(L_fake_A.shape).to(device)
D_ad_fake_loss_LA = Adversarial_loss(L_fake_A, fake_labels)
D_ad_loss_LA = D_ad_real_loss_LA + D_ad_fake_loss_LA
real_labels = torch.ones(L_real_A_cam.shape).to(device)
D_ad_cam_real_loss_LA = Adversarial_loss(L_real_A_cam, real_labels)
fake_labels = torch.zeros(L_fake_A_cam.shape).to(device)
D_ad_cam_fake_loss_LA = Adversarial_loss(L_fake_A_cam, fake_labels)
D_ad_cam_loss_LA = D_ad_cam_real_loss_LA + D_ad_cam_fake_loss_LA
real_labels = torch.ones(L_real_B.shape).to(device)
D_ad_real_loss_LB = Adversarial_loss(L_real_B, real_labels)
fake_labels = torch.zeros(L_fake_B.shape).to(device)
D_ad_fake_loss_LB = Adversarial_loss(L_fake_B, fake_labels)
D_ad_loss_LB = D_ad_real_loss_LB + D_ad_fake_loss_LB
real_labels = torch.ones(L_real_B_cam.shape).to(device)
D_ad_cam_real_loss_LB = Adversarial_loss(L_real_B_cam, real_labels)
fake_labels = torch.zeros(L_fake_B_cam.shape).to(device)
D_ad_cam_fake_loss_LB = Adversarial_loss(L_fake_B_cam, fake_labels)
D_ad_cam_loss_LB = D_ad_cam_real_loss_LB + D_ad_cam_fake_loss_LB
# Calculate Each Discriminator Loss #
D_loss_A = D_ad_loss_GA + D_ad_cam_loss_GA + D_ad_loss_LA + D_ad_cam_loss_LA
D_loss_B = D_ad_loss_GB + D_ad_cam_loss_GB + D_ad_loss_LB + D_ad_cam_loss_LB
# Calculate Total Discriminator Loss #
D_loss = D_loss_A + D_loss_B
# Back Propagation and Update #
D_loss.backward()
D_optim.step()
###################
# Train Generator #
###################
set_requires_grad([D_A, D_B, L_A, L_B], requires_grad=False)
# Forward Data #
fake_B, fake_B_cam, _ = G_A2B(real_A)
fake_A, fake_A_cam, _ = G_B2A(real_B)
fake_ABA, _, _ = G_B2A(fake_B)
fake_BAB, _, _ = G_A2B(fake_A)
fake_A2A, fake_A2A_cam, _ = G_A2B(real_A)
fake_B2B, fake_B2B_cam, _ = G_B2A(real_B)
G_fake_A, G_fake_A_cam, _ = D_A(fake_A)
L_fake_A, L_fake_A_cam, _ = L_A(fake_A)
G_fake_B, G_fake_B_cam, _ = D_B(fake_B)
L_fake_B, L_fake_B_cam, _ = L_B(fake_B)
# Adversarial Loss of Generator #
real_labels = torch.ones(G_fake_A.shape).to(device)
G_adv_fake_loss_A = Adversarial_loss(G_fake_A, real_labels)
real_labels = torch.ones(G_fake_A_cam.shape).to(device)
G_adv_cam_fake_loss_A = Adversarial_loss(G_fake_A_cam, real_labels)
G_adv_loss_A = G_adv_fake_loss_A + G_adv_cam_fake_loss_A
real_labels = torch.ones(G_fake_B.shape).to(device)
G_adv_fake_loss_B = Adversarial_loss(G_fake_B, real_labels)
real_labels = torch.ones(G_fake_B_cam.shape).to(device)
G_adv_cam_fake_loss_B = Adversarial_loss(G_fake_B_cam, real_labels)
G_adv_loss_B = G_adv_fake_loss_B + G_adv_cam_fake_loss_B
# Adversarial Loss of L #
real_labels = torch.ones(L_fake_A.shape).to(device)
L_adv_fake_loss_A = Adversarial_loss(L_fake_A, real_labels)
real_labels = torch.ones(L_fake_A_cam.shape).to(device)
L_adv_cam_fake_loss_A = Adversarial_loss(L_fake_A_cam, real_labels)
L_adv_loss_A = L_adv_fake_loss_A + L_adv_cam_fake_loss_A
real_labels = torch.ones(L_fake_B.shape).to(device)
L_adv_fake_loss_B = Adversarial_loss(L_fake_B, real_labels)
real_labels = torch.ones(L_fake_B_cam.shape).to(device)
L_adv_cam_fake_loss_B = Adversarial_loss(L_fake_B_cam, real_labels)
L_adv_loss_B = L_adv_fake_loss_B + L_adv_cam_fake_loss_B
# Cycle Consistency Loss #
G_recon_loss_A = Cycle_loss(fake_ABA, real_A)
G_recon_loss_B = Cycle_loss(fake_BAB, real_B)
G_identity_loss_A = Cycle_loss(fake_A2A, real_A)
G_identity_loss_B = Cycle_loss(fake_B2B, real_B)
G_cycle_loss_A = G_recon_loss_A + G_identity_loss_A
G_cycle_loss_B = G_recon_loss_B + G_identity_loss_B
# CAM Loss #
real_labels = torch.ones(fake_A_cam.shape).to(device)
G_cam_real_loss_A = BCE_loss(fake_A_cam, real_labels)
fake_labels = torch.zeros(fake_A2A_cam.shape).to(device)
G_cam_fake_loss_A = BCE_loss(fake_A2A_cam, fake_labels)
G_cam_loss_A = G_cam_real_loss_A + G_cam_fake_loss_A
real_labels = torch.ones(fake_B_cam.shape).to(device)
G_cam_real_loss_B = BCE_loss(fake_B_cam, real_labels)
fake_labels = torch.zeros(fake_B2B_cam.shape).to(device)
G_cam_fake_loss_B = BCE_loss(fake_B2B_cam, fake_labels)
G_cam_loss_B = G_cam_real_loss_B + G_cam_fake_loss_B
# Calculate Each Generator Loss #
G_loss_A = G_adv_loss_A + L_adv_loss_A + config.lambda_cycle * G_cycle_loss_A + config.lambda_cam * G_cam_loss_A
G_loss_B = G_adv_loss_B + L_adv_loss_B + config.lambda_cycle * G_cycle_loss_B + config.lambda_cam * G_cam_loss_B
# Calculate Total Generator Loss #
G_loss = G_loss_A + G_loss_B
# Back Propagation and Update #
G_loss.backward()
G_optim.step()
# Apply Rho Clipper to Generators #
G_A2B.apply(Rho_Clipper)
G_B2A.apply(Rho_Clipper)
# Add items to Lists #
D_losses.append(D_loss.item())
G_losses.append(G_loss.item())
####################
# Print Statistics #
####################
if (i + 1) % config.print_every == 0:
print(
"U-GAT-IT | Epochs [{}/{}] | Iterations [{}/{}] | D Loss {:.4f} | G Loss {:.4f}"
.format(epoch + 1, config.num_epochs, i + 1, total_batch,
np.average(D_losses), np.average(G_losses)))
# Save Sample Images #
save_samples(test_loader_selfie, G_A2B, epoch,
config.samples_path)
# Adjust Learning Rate #
D_optim_scheduler.step()
G_optim_scheduler.step()
# Save Model Weights #
if (epoch + 1) % config.save_every == 0:
torch.save(
D_A.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_D_A_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
D_B.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_D_B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
L_A.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_L_A_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
L_B.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_L_B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
G_A2B.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_G_A2B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
G_B2A.state_dict(),
os.path.join(config.weights_path,
'U-GAT-IT_G_B2A_Epoch_{}.pkl'.format(epoch + 1)))
# Plot Losses #
plot_losses(D_losses, G_losses, config.num_epochs, config.plots_path)
# Make a GIF file #
make_gifs_train('U-GAT-IT', config.samples_path)
print("Training finished.")
def train(dataset, dataset_folder, task, number_of_points, batch_size, epochs,
learning_rate, output_folder, number_of_workers, model_checkpoint):
train_dataset = DATASETS[dataset](dataset_folder,
task=task,
number_of_points=number_of_points)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=number_of_workers)
test_dataset = DATASETS[dataset](dataset_folder,
task=task,
train=False,
number_of_points=number_of_points)
test_dataloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=number_of_workers)
if task == 'classification':
model = ClassificationPointNet(
num_classes=train_dataset.NUM_CLASSIFICATION_CLASSES,
point_dimension=train_dataset.POINT_DIMENSION)
elif task == 'segmentation':
model = SegmentationPointNet(
num_classes=train_dataset.NUM_SEGMENTATION_CLASSES,
point_dimension=train_dataset.POINT_DIMENSION)
else:
raise Exception('Unknown task !')
if torch.cuda.is_available():
model.cuda()
if model_checkpoint:
model.load_state_dict(torch.load(model_checkpoint))
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
mb = master_bar(range(epochs))
if not os.path.isdir(output_folder):
os.mkdir(output_folder)
with open(os.path.join(output_folder, 'training_log.csv'), 'w+') as fid:
fid.write('train_loss,test_loss,train_accuracy,test_accuracy\n')
train_loss = []
test_loss = []
train_acc = []
test_acc = []
for epoch in mb:
epoch_train_loss = []
epoch_train_acc = []
batch_number = 0
for data in progress_bar(train_dataloader, parent=mb):
batch_number += 1
points, targets = data
# print(points.shape)
# print(points)
# print(targets)
# print(targets.shape)
if torch.cuda.is_available():
points, targets = points.cuda(), targets.cuda()
if points.shape[0] <= 1:
continue
optimizer.zero_grad()
model = model.train()
preds, feature_transform = model(points)
if task == 'segmentation':
preds = preds.view(-1, train_dataset.NUM_SEGMENTATION_CLASSES)
targets = targets.view(-1)
identity = torch.eye(feature_transform.shape[-1])
if torch.cuda.is_available():
identity = identity.cuda()
regularization_loss = torch.norm(identity - torch.bmm(
feature_transform, feature_transform.transpose(2, 1)))
loss = F.nll_loss(preds, targets) + 0.001 * regularization_loss
epoch_train_loss.append(loss.cpu().item())
loss.backward()
optimizer.step()
preds = preds.data.max(1)[1]
print(targets)
print(preds)
corrects = preds.eq(targets.data).cpu().sum()
if task == 'classification':
accuracy = corrects.item() / float(batch_size)
elif task == 'segmentation':
accuracy = corrects.item() / float(
batch_size * number_of_points)
epoch_train_acc.append(accuracy)
mb.child.comment = 'train loss: %f, train accuracy: %f' % (
np.mean(epoch_train_loss), np.mean(epoch_train_acc))
epoch_test_loss = []
epoch_test_acc = []
for batch_number, data in enumerate(test_dataloader):
points, targets = data
if torch.cuda.is_available():
points, targets = points.cuda(), targets.cuda()
model = model.eval()
preds, feature_transform = model(points)
if task == 'segmentation':
preds = preds.view(-1, train_dataset.NUM_SEGMENTATION_CLASSES)
targets = targets.view(-1)
loss = F.nll_loss(preds, targets)
epoch_test_loss.append(loss.cpu().item())
preds = preds.data.max(1)[1]
corrects = preds.eq(targets.data).cpu().sum()
if task == 'classification':
accuracy = corrects.item() / float(batch_size)
elif task == 'segmentation':
accuracy = corrects.item() / float(
batch_size * number_of_points)
epoch_test_acc.append(accuracy)
mb.write(
'Epoch %s: train loss: %s, val loss: %f, train accuracy: %s, val accuracy: %f'
% (epoch, np.mean(epoch_train_loss), np.mean(epoch_test_loss),
np.mean(epoch_train_acc), np.mean(epoch_test_acc)))
if test_acc and np.mean(epoch_test_acc) > np.max(test_acc):
torch.save(
model.state_dict(),
os.path.join(output_folder, 'shapenet_%s_model.pth' % task))
with open(os.path.join(output_folder, 'training_log.csv'), 'a') as fid:
fid.write(
'%s,%s,%s,%s,%s\n' %
(epoch, np.mean(epoch_train_loss), np.mean(epoch_test_loss),
np.mean(epoch_train_acc), np.mean(epoch_test_acc)))
train_loss.append(np.mean(epoch_train_loss))
test_loss.append(np.mean(epoch_test_loss))
train_acc.append(np.mean(epoch_train_acc))
test_acc.append(np.mean(epoch_test_acc))
plot_losses(train_loss,
test_loss,
save_to_file=os.path.join(output_folder, 'loss_plot.png'))
plot_accuracies(train_acc,
test_acc,
save_to_file=os.path.join(output_folder,
'accuracy_plot.png'))
def loop(data_loader,
num_epochs=1000,
save_every=1000,
train_losses=[],
test_losses=[],
train_cnts=[],
test_cnts=[],
dummy=False):
print("starting training loop for data with %s batches" %
data_loader.num_batches)
st = time.time()
if len(train_losses):
# resume cnt from last save
last_save = train_cnts[-1]
cnt = train_cnts[-1]
else:
last_save = 0
cnt = 0
for e in range(num_epochs):
ecnt = 0
tst = round((time.time() - st) / 60., 0)
if not e % 1 and e > 0:
print(
"starting epoch %s, %s mins, loss %s, seen %s, last save at %s"
% (e, tst, train_losses[-1], cnt, last_save))
batch_loss = []
for b in range(data_loader.num_batches):
xnp, ynp = data_loader.next_batch()
x = Variable(torch.FloatTensor(xnp))
y = Variable(torch.FloatTensor(ynp))
y_pred, loss = train(x, y, validation=False)
train_cnts.append(cnt)
train_losses.append(loss)
if cnt % 100:
valy_pred, val_mean_loss = train(v_x, v_y, validation=True)
test_losses.append(val_mean_loss)
test_cnts.append(cnt)
if cnt - last_save >= save_every:
last_save = cnt
# find test loss
print(
'epoch: {} saving after example {} train loss {} test loss {}'
.format(e, cnt, loss, val_mean_loss))
state = {
'train_cnts': train_cnts,
'train_losses': train_losses,
'test_cnts': test_cnts,
'test_losses': test_losses,
'state_dict': lstm.state_dict(),
'optimizer': optim.state_dict(),
}
basename = os.path.join(savedir,
'%s_%015d' % (model_save_name, cnt))
n = 500
plot_losses(rolling_average(train_cnts, n),
rolling_average(train_losses, n),
rolling_average(test_cnts, n),
rolling_average(test_losses, n),
name=basename + '_loss.png')
save_checkpoint(state, filename=basename + '.pkl')
cnt += x.shape[1]
ecnt += x.shape[1]
def train():
# Fix Seed for Reproducibility #
torch.manual_seed(9)
if torch.cuda.is_available():
torch.cuda.manual_seed(9)
# Samples, Weights and Results Path #
paths = [config.samples_path, config.weights_path, config.plots_path]
paths = [make_dirs(path) for path in paths]
# Prepare Data Loader #
train_horse_loader, train_zebra_loader = get_horse2zebra_loader(
purpose='train', batch_size=config.batch_size)
test_horse_loader, test_zebra_loader = get_horse2zebra_loader(
purpose='test', batch_size=config.val_batch_size)
total_batch = min(len(train_horse_loader), len(train_zebra_loader))
# Prepare Networks #
D_A = Discriminator()
D_B = Discriminator()
G_A2B = Generator()
G_B2A = Generator()
networks = [D_A, D_B, G_A2B, G_B2A]
for network in networks:
network.to(device)
# Loss Function #
criterion_Adversarial = nn.MSELoss()
criterion_Cycle = nn.L1Loss()
criterion_Identity = nn.L1Loss()
# Optimizers #
D_A_optim = torch.optim.Adam(D_A.parameters(),
lr=config.lr,
betas=(0.5, 0.999))
D_B_optim = torch.optim.Adam(D_B.parameters(),
lr=config.lr,
betas=(0.5, 0.999))
G_optim = torch.optim.Adam(chain(G_A2B.parameters(), G_B2A.parameters()),
lr=config.lr,
betas=(0.5, 0.999))
D_A_optim_scheduler = get_lr_scheduler(D_A_optim)
D_B_optim_scheduler = get_lr_scheduler(D_B_optim)
G_optim_scheduler = get_lr_scheduler(G_optim)
# Lists #
D_losses_A, D_losses_B, G_losses = [], [], []
# Training #
print("Training CycleGAN started with total epoch of {}.".format(
config.num_epochs))
for epoch in range(config.num_epochs):
for i, (horse,
zebra) in enumerate(zip(train_horse_loader,
train_zebra_loader)):
# Data Preparation #
real_A = horse.to(device)
real_B = zebra.to(device)
# Initialize Optimizers #
G_optim.zero_grad()
D_A_optim.zero_grad()
D_B_optim.zero_grad()
###################
# Train Generator #
###################
set_requires_grad([D_A, D_B], requires_grad=False)
# Adversarial Loss #
fake_A = G_B2A(real_B)
prob_fake_A = D_A(fake_A)
real_labels = torch.ones(prob_fake_A.size()).to(device)
G_mse_loss_B2A = criterion_Adversarial(prob_fake_A, real_labels)
fake_B = G_A2B(real_A)
prob_fake_B = D_B(fake_B)
real_labels = torch.ones(prob_fake_B.size()).to(device)
G_mse_loss_A2B = criterion_Adversarial(prob_fake_B, real_labels)
# Identity Loss #
identity_A = G_B2A(real_A)
G_identity_loss_A = config.lambda_identity * criterion_Identity(
identity_A, real_A)
identity_B = G_A2B(real_B)
G_identity_loss_B = config.lambda_identity * criterion_Identity(
identity_B, real_B)
# Cycle Loss #
reconstructed_A = G_B2A(fake_B)
G_cycle_loss_ABA = config.lambda_cycle * criterion_Cycle(
reconstructed_A, real_A)
reconstructed_B = G_A2B(fake_A)
G_cycle_loss_BAB = config.lambda_cycle * criterion_Cycle(
reconstructed_B, real_B)
# Calculate Total Generator Loss #
G_loss = G_mse_loss_B2A + G_mse_loss_A2B + G_identity_loss_A + G_identity_loss_B + G_cycle_loss_ABA + G_cycle_loss_BAB
# Back Propagation and Update #
G_loss.backward(retain_graph=True)
G_optim.step()
#######################
# Train Discriminator #
#######################
set_requires_grad([D_A, D_B], requires_grad=True)
## Train Discriminator A ##
# Real Loss #
prob_real_A = D_A(real_A)
real_labels = torch.ones(prob_real_A.size()).to(device)
D_real_loss_A = criterion_Adversarial(prob_real_A, real_labels)
# Fake Loss #
prob_fake_A = D_A(fake_A.detach())
fake_labels = torch.zeros(prob_fake_A.size()).to(device)
D_fake_loss_A = criterion_Adversarial(prob_fake_A, fake_labels)
# Calculate Total Discriminator A Loss #
D_loss_A = config.lambda_identity * (D_real_loss_A +
D_fake_loss_A).mean()
# Back propagation and Update #
D_loss_A.backward(retain_graph=True)
D_A_optim.step()
## Train Discriminator B ##
# Real Loss #
prob_real_B = D_B(real_B)
real_labels = torch.ones(prob_real_B.size()).to(device)
loss_real_B = criterion_Adversarial(prob_real_B, real_labels)
# Fake Loss #
prob_fake_B = D_B(fake_B.detach())
fake_labels = torch.zeros(prob_fake_B.size()).to(device)
loss_fake_B = criterion_Adversarial(prob_fake_B, fake_labels)
# Calculate Total Discriminator B Loss #
D_loss_B = config.lambda_identity * (loss_real_B +
loss_fake_B).mean()
# Back propagation and Update #
D_loss_B.backward(retain_graph=True)
D_B_optim.step()
# Add items to Lists #
D_losses_A.append(D_loss_A.item())
D_losses_B.append(D_loss_B.item())
G_losses.append(G_loss.item())
####################
# Print Statistics #
####################
if (i + 1) % config.print_every == 0:
print(
"CycleGAN | Epoch [{}/{}] | Iterations [{}/{}] | D_A Loss {:.4f} | D_B Loss {:.4f} | G Loss {:.4f}"
.format(epoch + 1, config.num_epochs, i + 1, total_batch,
np.average(D_losses_A), np.average(D_losses_B),
np.average(G_losses)))
# Save Sample Images #
sample_images(test_horse_loader, test_zebra_loader, G_A2B,
G_B2A, epoch, config.samples_path)
# Adjust Learning Rate #
D_A_optim_scheduler.step()
D_B_optim_scheduler.step()
G_optim_scheduler.step()
# Save Model Weights #
if (epoch + 1) % config.save_every == 0:
torch.save(
G_A2B.state_dict(),
os.path.join(
config.weights_path,
'CycleGAN_Generator_A2B_Epoch_{}.pkl'.format(epoch + 1)))
torch.save(
G_B2A.state_dict(),
os.path.join(
config.weights_path,
'CycleGAN_Generator_B2A_Epoch_{}.pkl'.format(epoch + 1)))
# Make a GIF file #
make_gifs_train("CycleGAN", config.samples_path)
# Plot Losses #
plot_losses(D_losses_A, D_losses_B, G_losses, config.num_epochs,
config.plots_path)
print("Training finished.")
use_test=False)
# train_leaf_starting_weights = list()
# for i in range(3):
# train_leaf_starting_weights.append(model_leaves[2].weight_list[i].detach().numpy())
delta_losses, prediction_losses, validation_losses, _ = experiment.train_dendronet(
)
simple_prediction_losses, simple_validation_losses, _ = experiment.train_simple_model(
)
# todo: omitted a block of code that collects that targets and final predictions for every model at every seed
train_leaves = np.take(leaves, experiment.train_idx)
train_model_leaves = np.take(model_leaves, experiment.train_idx)
# plotting losses
x_label = str(config['validation_interval']) + "s of steps"
plot_file = os.path.join(output_dir, 'delta_loss.png')
plot_losses(plot_file, [delta_losses], ['delta_loss'], x_label)
plot_file = os.path.join(output_dir, 'dendronet_prediction_losses.png')
plot_losses(plot_file, [prediction_losses, validation_losses],
['training', 'validation'], x_label)
plot_file = os.path.join(output_dir, 'simple_losses.png')
plot_losses(plot_file,
[simple_prediction_losses, simple_validation_losses],
['training', 'validation'], x_label)
plot_file = os.path.join(output_dir, 'validation_loss_comparison.png')
plot_losses(plot_file, [validation_losses, simple_validation_losses],
['dendro_valid', 'baseline_valid'], x_label)
# todo: omitted analysis of final delta loss vs ideal
"""
Retrieving weights and analysing disentanglement metrics for simple model
For now sticking with the assumption that we have an diagonal matrix style generative process
We do not know the configuration of factors->rows, so we charitably assume that the best possible
def train_mode(gen,
dis,
trainLoader,
useNoise=False,
beta1=0.5,
c=0.01,
k=1,
WGAN=False):
####### Define optimizer #######
genOptimizer = optim.Adam(gen.parameters(),
lr=opts.lr,
betas=(beta1, 0.999))
disOptimizer = optim.Adam(dis.parameters(),
lr=opts.lr,
betas=(beta1, 0.999))
if gen.useCUDA:
torch.cuda.set_device(opts.gpuNo)
gen.cuda()
dis.cuda()
####### Create a new folder to save results and model info #######
exDir = make_new_folder(opts.outDir)
print 'Outputs will be saved to:', exDir
save_input_args(exDir, opts)
#noise level
noiseSigma = np.logspace(np.log2(0.5),
np.log2(0.001),
opts.maxEpochs,
base=2)
####### Start Training #######
losses = {'gen': [], 'dis': []}
for e in range(opts.maxEpochs):
dis.train()
gen.train()
epochLoss_gen = 0
epochLoss_dis = 0
noiseLevel = float(noiseSigma[e])
T = time()
for i, data in enumerate(trainLoader, 0):
for _ in range(k):
# add a small amount of corruption to the data
xReal = Variable(data[0])
if gen.useCUDA:
xReal = xReal.cuda()
if useNoise:
xReal = corrupt(xReal, noiseLevel) #add a little noise
####### Calculate discriminator loss #######
noSamples = xReal.size(0)
xFake = gen.sample_x(noSamples)
if useNoise:
xFake = corrupt(xFake, noiseLevel) #add a little noise
pReal_D = dis.forward(xReal)
pFake_D = dis.forward(xFake.detach())
real = dis.ones(xReal.size(0))
fake = dis.zeros(xFake.size(0))
if WGAN:
disLoss = pFake_D.mean() - pReal_D.mean()
else:
disLoss = opts.pi * F.binary_cross_entropy(pReal_D, real) + \
(1 - opts.pi) * F.binary_cross_entropy(pFake_D, fake)
####### Do DIS updates #######
disOptimizer.zero_grad()
disLoss.backward()
disOptimizer.step()
#### clip DIS weights #### YM
if WGAN:
for p in dis.parameters():
p.data.clamp_(-c, c)
losses['dis'].append(disLoss.data[0])
####### Calculate generator loss #######
xFake_ = gen.sample_x(noSamples)
if useNoise:
xFake_ = corrupt(xFake_, noiseLevel) #add a little noise
pFake_G = dis.forward(xFake_)
if WGAN:
genLoss = -pFake_G.mean()
else:
genLoss = F.binary_cross_entropy(pFake_G, real)
####### Do GEN updates #######
genOptimizer.zero_grad()
genLoss.backward()
genOptimizer.step()
losses['gen'].append(genLoss.data[0])
####### Print info #######
if i % 100 == 1:
print '[%d, %d] gen: %.5f, dis: %.5f, time: %.2f' \
% (e, i, genLoss.data[0], disLoss.data[0], time()-T)
####### Tests #######
gen.eval()
print 'Outputs will be saved to:', exDir
#save some samples
samples = gen.sample_x(49)
save_image(samples.data,
join(exDir, 'epoch' + str(e) + '.png'),
normalize=True)
#plot
plot_losses(losses, exDir, epochs=e + 1)
####### Save params #######
gen.save_params(exDir)
dis.save_params(exDir)
return gen, dis
def train(params):
'''
Train the GAN network.
'''
criterion = nn.BCEWithLogitsLoss()
cur_step = 0
mean_generator_loss = 0
mean_discriminator_loss = 0
dataloader = get_train_data(params)
display_step = 500
g_loss = []
d_loss = []
gen, gen_opt = create_generator(params)
disc, disc_opt = create_discriminator(params)
for epoch in range(params['n_epochs']):
# Dataloader returns the batches
for real, _ in tqdm(dataloader):
cur_batch_size = len(real)
real = real.to(params['device'])
## Update discriminator ##
disc_opt.zero_grad()
fake_noise = get_noise(cur_batch_size, params)
fake = gen(fake_noise)
disc_fake_pred = disc(fake.detach())
# pylint: disable=E1101
disc_fake_loss = criterion(disc_fake_pred, torch.zeros_like(disc_fake_pred))
disc_real_pred = disc(real)
disc_real_loss = criterion(disc_real_pred, torch.ones_like(disc_real_pred))
disc_loss = (disc_fake_loss + disc_real_loss) / 2
# pylint: enable=E1101
# Keep track of the average discriminator loss
mean_discriminator_loss += disc_loss.item() / display_step
# Update gradients
disc_loss.backward(retain_graph=True)
# Update optimizer
disc_opt.step()
## Update generator ##
gen_opt.zero_grad()
fake_noise_2 = get_noise(cur_batch_size, params)
fake_2 = gen(fake_noise_2)
disc_fake_pred = disc(fake_2)
# pylint: disable=E1101
gen_loss = criterion(disc_fake_pred, torch.ones_like(disc_fake_pred))
# pylint: enable=E1101
gen_loss.backward()
gen_opt.step()
# Keep track of the average generator loss
mean_generator_loss += gen_loss.item() / display_step
g_loss.append(gen_loss.item())
d_loss.append(disc_loss.item())
## Visualization code ##
if cur_step % display_step == 0 and cur_step > 0:
print(f"Step {cur_step}: Generator loss: {mean_generator_loss}, discriminator loss: {mean_discriminator_loss}")
show_tensor_images(fake)
show_tensor_images(real)
mean_generator_loss = 0
mean_discriminator_loss = 0
cur_step += 1
plot_losses(g_loss, d_loss)
# Save the trained model.
torch.save({
'generator' : gen.state_dict(),
'discriminator' : disc.state_dict(),
'optimizerG' : gen_opt.state_dict(),
'optimizerD' : disc_opt.state_dict(),
'params' : params
}, 'model/model_final.pth')
config,
data_root,
leaves,
baselines=args.baselines,
expanded_x=None,
use_test=True)
par_dendro_losses, par_prediction_losses, par_validation_losses, par_validation_aucs = baseline_experiment.train_dendronet(
)
dump_dict(config, output_dir)
x_label = str(config['validation_interval']) + "s of steps"
if args.baselines:
plot_file = os.path.join(output_dir, 'parsimony_losses.png')
plot_losses(plot_file, [
par_dendro_losses, par_prediction_losses, par_validation_losses
], ['mutation', 'training', 'validation'], x_label)
if len(par_validation_aucs) >= 1:
if args.baselines:
aucs['parsimony_best'].append(max(par_validation_aucs[1:]))
aucs['parsimony_final'].append(par_validation_aucs[-1])
auc_list = list()
log_auc_names = list()
for name in basenames:
auc_list.append(aucs[name])
log_auc_names.append(name)
if len(auc_list[0]) > 0:
log_aucs(os.path.join(base_output_dir, 'auc_log.txt'), auc_list,
log_auc_names)
simple_predictions = [simple_model.call(leaf.x).numpy() for leaf in leaves]
fig_targets = [leaf.y for leaf in leaves]
fig_species_names = [leaf.name for leaf in leaves]
fig_dict = {
'dendro_predictions': dendro_predictions,
'simple_predictions': simple_predictions,
'true_classifications': fig_targets,
'species_names': fig_species_names
}
dump_dict(fig_dict, output_dir, name=tree_labels_file_name)
# print_tree_model(tree_model.root_layer, root=data_root, lifestyle=ls, feature_index=FEATURE_INDEX)
x_label = str(config['validation_interval']) + "s of steps"
plot_file = os.path.join(output_dir, 'dendronet_losses.png')
plot_losses(plot_file, [dendronet_losses, prediction_losses, validation_losses],
['mutation', 'training', 'validation'], x_label)
plot_file = os.path.join(output_dir, 'simple_losses.png')
plot_losses(plot_file, [simple_prediction_losses, simple_validation_losses], ['training', 'validation'], x_label)
if args.baselines:
plot_file = os.path.join(output_dir, 'parsimony_losses.png')
plot_losses(plot_file, [par_dendro_losses, par_prediction_losses, par_validation_losses],
['mutation', 'training', 'validation'], x_label)
plot_file = os.path.join(output_dir, 'one_hot_losses.png')
plot_losses(plot_file, [one_hot_prediction_losses, one_hot_validation_losses], ['training', 'validation'], x_label)
if len(simple_validation_aucs) >= 1:
aucs[ls]['dendro_best'].append(max(validation_aucs[1:]))
aucs[ls]['simple_best'].append(max(simple_validation_aucs[1:]))
aucs[ls]['dendro_final'].append(validation_aucs[-1])
aucs[ls]['simple_final'].append(simple_validation_aucs[-1])
gnn_model.to(device)
(trained_gnn_model, train_loss_gnn, val_loss_gnn) = train_model(gnn_model,
train_data,
val_data,
n_epochs=3)
# %%
print(f"train_loss_gnn: {train_loss_gnn}")
print(f"val_loss_gnn: {val_loss_gnn}")
# This doesn't work
# gnn_model.save("gnn40epochs.pth")
plt.rcParams.update({"text.usetex": False})
plot_losses("GNN model", train_loss_gnn, val_loss_gnn)
# %% [markdown]
## Probemos con más cervezas a la vez
### Elegimos cervezas con muchas reviews
# %%
n_target_beers = 1
n_ratings_by_beer = (~df_user_rating_train.isnull()).sum()
random_beers = (n_ratings_by_beer.sort_values(
ascending=False).iloc[:200].sample(n_target_beers).index)
random_beer_names = df_beer.loc[random_beers.astype(int), "beer_name"]
random_beers = [str(beer_id) for beer_id in random_beers]
print(random_beer_names)
