def compare_VAE(batch, generator, epoch, path):
'''
Saves parallely generated batch/reconstructed batch
images next to each other.
Keyword arguments:
> batch (torch.tensor) -- original input data.
> generator (torch.nn.Module) -- current generator.
> epoch (int) -- number of COMPLETED training epochs.
> path (string) -- save path.
Returns: N/A
'''
print('Generating comparison images from VAE...')
if not os.path.isdir(path):
os.makedirs(path)
save_path = path + '/reconstruction_' + str(
epoch) + '~' + constants.get_cur_time() + '.png'
with torch.no_grad():
n = min(batch.size(0), MAX_IMG)
_, recon_batch, _, _ = generator(batch)
comparison = torch.cat(
[batch[:n],
recon_batch.view(batch.size(0), 1, 28, 28)[:n]])
save_and_upload_image(comparison.cpu(), save_path, nrow=n)
print('Finished! Samples saved under ' + save_path)
def convert():
'''
Converts all old filenames into timestamped ones (locally).
SHOULD ONLY BE USED ONCE!!!
'''
for folder_name in SYNC_DIRECTORIES:
convert_helper(folder_name, constants.get_cur_time())
def visualize_attack(params, path):
'''
Saves the image of the raw data from a dataloader,
alongside its perturbed version.
'''
data, perturbed_data = train_utils.sample_attack_from_dataset(params)
print(data.min(), data.max(), perturbed_data.min(), perturbed_data.max())
save_and_upload_image(
data.cpu(), path + '_regular~' + constants.get_cur_time() + '.png')
save_and_upload_image(
perturbed_data.cpu(),
path + '_attack~' + constants.get_cur_time() + '.png')
save_and_upload_image(data.cpu(),
path + '_regular_norm~' + constants.get_cur_time() +
'.png',
normalize=True)
save_and_upload_image(perturbed_data.cpu(),
path + '_attack_norm~' + constants.get_cur_time() +
'.png',
normalize=True)
print('Finished visualizing!')
def general_plot(plot_points,
plot_names,
run_name,
title,
xlabel,
ylabel,
plt_symbol='-o',
figsize=(10, 8),
scales=[1]):
'''
Generic plotting function to be used for acc/loss.
Keyword arguments:
> plot_points (list[list]) -- list of point lists to be plotted.
> plot_names (list[string]) -- parallel list of plot names.
> run_name (string) -- current run to save under.
> title (string) -- plot title (also saved extension).
> xlabel (string) -- plot x-axis label
> ylabel (string) -- plot y-axis label
'''
save_name = 'graphs/' + run_name + '/' + run_name + '_' + title.lower() + \
'~' + constants.get_cur_time() + '.png'
if not os.path.isdir('graphs/' + run_name + '/'):
os.makedirs('graphs/' + run_name + '/')
plt.figure(figsize=figsize)
plt.title(title)
for idx, plot in enumerate(plot_points):
x_labels = list(i * scales[idx] for i in range(len(plot)))
plt.plot(x_labels, plot, plt_symbol, label=plot_names[idx])
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.legend(loc='lower right')
plt.savefig(save_name)
plt.close()
# Syncs with Box
rid_duplicate(save_name)
box_utils.upload_single(save_name)
def sample_VAE(vae_model, device, epoch, path):
'''
Returns randomly generated samples from vae_model.
Keyword arguments:
> vae_model (torch.nn.Module) -- generator model.
> device (torch.device) -- CUDA or cpu.
> epoch (int), path (string) -- for save path.
Returns: N/A
'''
print('Sampling from VAE...')
if not os.path.isdir(path):
os.makedirs(path)
save_path = path + '/sample_' + str(
epoch) + '~' + constants.get_cur_time() + '.png'
vae_model.eval()
with torch.no_grad():
sample = torch.randn(64, 128).to(device)
sample = vae_model.decode(sample).cpu()
save_and_upload_image(sample.view(64, 1, 28, 28), save_path)
print('Finished! Samples saved under ' + save_path)
def save_checkpoint(params, epoch):
'''
Saves the parameter dictionary passed in.
`epoch` is assumed to be the number of training
epochs completed by the model (i.e. begin
training at epoch + 1).
Keyword arguments:
> params (dict) -- current state variable
> epoch (int) -- number of COMPLETED training epochs
Returns: N/A
'''
save_path = 'models/' + model_type(params) + '/' + params['run_name'] \
+ '/' + params['run_name'] + '_epoch_' + str(epoch) + \
'~' + constants.get_cur_time() + DEFAULT_SAVE_EXTENSION
print('Saving current state to', save_path)
torch.save(params, save_path)
# Uploads single saved checkpoint to Box
viz_utils.rid_duplicate(save_path)
box_utils.upload_single(save_path)