def test_disent(cfg):
print("Testing image disentanglement.")
# load the data
data, loader = get_mnist_dataset(cfg, 'disent')
# load the model and set criterion
models = load_disent_models(cfg)
for name, model in models.items():
model.eval()
# get the data
x = next(iter(loader.tr))[0]
pics = [x_i.to(cfg.disent.device) for x_i in x]
nt = len(pics)
fig, ax = plt.subplots(nt, nt + 1, figsize=(8, 8))
for i, pic_i in enumerate(pics):
plot_th_tensor(ax, i, -1, pic_i)
encC_i = models.enc_c(pic_i)
encD_i = models.enc_d(pic_i)
for j, pic_j in enumerate(pics):
encC_j = models.enc_c(pic_j)
encD_j = models.enc_d(pic_j)
dec_ij = models.dec(torch.cat([encC_i, encD_j], dim=1))
plot_th_tensor(ax, i, j, dec_ij)
if i == j: continue
dec_ji = models.dec(torch.cat([encC_j, encD_i], dim=1))
# plot_th_tensor(ax,j,i,dec_ji)
plt.savefig(f"test_disentangle_{cfg.disent.epoch_num}.png")
plt.clf()
plt.cla()
def exploring_nt_xent_loss(cfg):
print("Exploring the NT_Xent loss function.")
# load the data
data, loader = get_mnist_dataset(cfg, 'disent')
img_set, img = next(iter(loader.tr))
print(len(img_set))
print(len(img_set[0]))
print(len(img))
print(len(img[0]))
# x = [x_i.to('cpu')[0].permute((1,2,0)).numpy() for x_i in x]
# x_dec = [x_dec_i.to('cpu')[0].detach().permute((1,2,0)).numpy() for x_dec_i in x_dec]
batch_size = cfg.disent.batch_size
loss = NT_Xent(batch_size, 0.1, cfg.disent.device, 1)
print(img.reshape(batch_size, -1).shape)
_, img1 = next(iter(loader.tr))
_, img2 = next(iter(loader.tr))
img1 = img1.reshape(batch_size, -1)
img2 = img2.reshape(batch_size, -1)
print(loss(img1, img2))
fig, ax = plt.subplots(len(img_set), 2)
for i, img_i in enumerate(img_set):
# x_recon = x_dec_i - np.min(x_dec_i)
# x_recon = x_recon / np.max(x_recon)
# print(x_noisy.shape,x_recon.shape)
ax[i, 0].imshow(img_i[0])
ax[i, 1].imshow(img[0])
plt.savefig("mnist_rand_blocks.png")
plt.clf()
plt.cla()
def train_disent_exp(cfg):
print("Training disentangled representations.")
# load the data
data, loader = get_mnist_dataset(cfg, 'disent')
# load the model and set criterion
models = load_disent_models(cfg)
hyperparams = edict()
hyperparams.h = 0.5
hyperparams.g = 0.5
hyperparams.x = 0.5
hyperparams.temperature = 0.1
criterion = DisentangleLoss(models, hyperparams, cfg.disent.N,
cfg.disent.batch_size, cfg.disent.device)
optimizer, scheduler = get_disent_optim(models)
# init the training loop
writer = SummaryWriter()
global_step, current_epoch = get_model_epoch_info(cfg.disent)
cfg.disent.global_step = global_step
cfg.disent.current_epoch = current_epoch
# training loop
for epoch in range(cfg.disent.current_epoch, cfg.disent.epochs + 1):
lr = optimizer.param_groups[0]["lr"]
loss_epoch = train_disent(cfg.disent, loader.tr, models, criterion,
optimizer, epoch, writer)
if scheduler:
scheduler.step(loss_epoch)
if epoch % cfg.disent.checkpoint_interval == 0:
save_disent_models(cfg, models, optimizer)
writer.add_scalar("Loss/train", loss_epoch / len(loader.tr), epoch)
writer.add_scalar("Misc/learning_rate", lr, epoch)
print(
f"Epoch [{epoch}/{cfg.disent.epochs}]\t Loss: {loss_epoch / len(loader.tr)}\t lr: {round(lr, 5)}"
)
cfg.disent.current_epoch += 1
save_disent_models(cfg, models, optimizer)
import torch
from torch import optim
from datasets import get_mnist_dataset, get_cifar10_dataset, get_data_loader
from utils import *
from models import *
trainset, testset = get_mnist_dataset()
trainloader, testloader = get_data_loader(trainset, testset)
batch, labels = next(iter(trainloader))
plot_batch(batch)
batch_var = Variable(batch.cuda())
labels_var = Variable(one_hotify(labels).cuda())
base_model = BaselineCNN().cuda()
print(count_params(base_model))
base_loss = nn.CrossEntropyLoss()
base_optimizer = optim.Adam(base_model.parameters())
base_trainer = Trainer(base_model,
base_optimizer,
base_loss,
trainloader,
testloader,
use_cuda=True)
base_trainer.run(epochs=10)
base_trainer.save_checkpoint('weights/baseline_mnist.pth.tar')