def plot_centroid(args, model, data_loader, mus, filename=None, n_clusters=10):
if filename is None:
kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(mus)
filename = args.dataset + "-kmeans-" + str(args.latent_dim) + "-beta-" + str(args.beta) + ".pt"
path = os.path.join(args.save_dir, filename)
print("Saving kmeans: %s\n" % (path))
pickle.dump(kmeans, open(path, "wb"))
else:
#kmeans = KMeans(n_clusters=n_clusters, random_state=0)
path = os.path.join(args.save_dir, filename)
print("Loading kmeans: %s\n" % (path))
kmeans = pickle.load(open(path, "rb"))
centers = kmeans.cluster_centers_
model.eval()
device = get_device()
with torch.no_grad():
ytrue = torch.tensor(np.array([]), dtype=torch.long)
ypred = torch.tensor(np.array([]), dtype=torch.long)
ytrue = ytrue.to(device)
ypred = ypred.to(device)
for i, (x, target) in enumerate(data_loader):
x = x.to(device)
# print('x shape', x.shape)
target = target.to(device)
mu, logvar = model.encoder(x)
labels = kmeans.predict(mu.cpu())
#print('labels shape', labels.shape)
#y = to_categorical(torch.from_numpy(labels).to(device, dtype=torch.long), n_clusters=n_clusters)
y = torch.from_numpy(labels).to(device, dtype=torch.long)
ytrue = torch.cat((ytrue, target), axis=0)
ypred = torch.cat((ypred, y.to(device)), axis=0)
accuracy = unsupervised_labels(ytrue.cpu(), ypred.cpu(), n_clusters, n_clusters)
method = "Unsupervised (VAE)"
progress_bar(i,
len(data_loader),
'%s | Acc: %0.2f%%'
% (method, accuracy))
if i==0:
centroid_mu = centers[labels]
centroid_mu = torch.tensor(centroid_mu).to(device)
# print('mu and logvar shapes', centroid_mu.shape, logvar.shape)
z = model.reparameterize(centroid_mu, logvar)
centroid = model.decoder(z)
n = min(x.size(0), 64)
folder = "results"
os.makedirs(folder, exist_ok=True)
xp, mu, logvar = model(x)
img = torch.cat([x[:n],
xp.view(args.batch_size, 1, args.crop_size, args.crop_size)[:n]])
filename = folder + '/recon.png'
save_image(img.cpu(), filename, nrow=n)
img = torch.cat([x[:n],
centroid.view(args.batch_size, 1, args.crop_size, args.crop_size)[:n]])
filename = folder + '/centroid.png'
save_image(img.cpu(), filename, nrow=n)
def test(args, model, test_loader, epoch=0):
model.eval()
recon_losses = AverageMeter()
kl_losses = AverageMeter()
device = get_device()
with torch.no_grad():
for i, (x, _) in enumerate(test_loader):
x = x.to(device)
mu, logvar = model.encoder(x)
noise = torch.randn_like(mu) * 2 * logvar.exp().pow(0.5)
noised_mu = mu + noise
z = model.reparameterize(noised_mu, logvar)
noised_pos = model.decoder(z)
noise = torch.randn_like(mu) * 2 * logvar.exp().pow(0.5)
noised_mu = mu + noise
z = model.reparameterize(noised_mu, logvar)
noised_neg = model.decoder(z)
xp, mu, logvar = model(x)
loss, recon, kl = loss_function(args, xp, x, mu, logvar, args.beta)
recon_losses.update(recon.float().item() / x.size(0))
kl_losses.update(kl.float().item() / x.size(0))
progress_bar(
i, len(test_loader),
'Epoch: %d | Recon Loss: %.6f | KL Loss: %.6f' %
(epoch, recon_losses.avg, kl_losses.avg))
if i == 0:
n = min(x.size(0), 64)
folder = "results"
os.makedirs(folder, exist_ok=True)
img = torch.cat([
x[:n],
xp.view(args.batch_size, 1, args.crop_size,
args.crop_size)[:n]
])
filename = folder + '/recon.png'
save_image(img.cpu(), filename, nrow=n)
img = torch.cat([
x[:n],
noised_pos.view(args.batch_size, 1, args.crop_size,
args.crop_size)[:n]
])
filename = folder + '/noised_pos.png'
save_image(img.cpu(), filename, nrow=n)
img = torch.cat([
x[:n],
noised_neg.view(args.batch_size, 1, args.crop_size,
args.crop_size)[:n]
])
filename = folder + '/noised_neg.png'
save_image(img.cpu(), filename, nrow=n)
def train(args, model, train_loader, optimizer, epoch):
model.train()
recon_losses = AverageMeter()
kl_losses = AverageMeter()
device = get_device()
for i, (x, _) in enumerate(train_loader):
x = x.to(device)
optimizer.zero_grad()
xp, mu, logvar = model(x)
loss, recon, kl = loss_function(args, xp, x, mu, logvar, args.beta)
loss.backward()
recon_losses.update(recon.float().item()/x.size(0))
kl_losses.update(kl.float().item()/x.size(0))
optimizer.step()
progress_bar(i,
len(train_loader),
'Epoch: %d | Recon Loss: %.6f | KL Loss: %.6f' % (epoch, recon_losses.avg, kl_losses.avg))
def tsne(args, model, data_loader, tsne=False):
model.eval()
device = get_device()
mus = np.array([])
with torch.no_grad():
for i, (x, _) in enumerate(data_loader):
x = x.to(device)
mu, logvar = model.encoder(x)
if len(mus) == 0:
mus = mu.cpu()
else:
mus = np.concatenate((mus, mu.cpu()), axis=0)
progress_bar(i, len(data_loader))
if tsne:
mus_tsne = TSNE(n_components=2).fit_transform(mus)
print(mus_tsne.shape)
return mus, mus_tsne
print(mus.shape)
return mus, _