File size: 1,883 Bytes
cb9235c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import torch
import numpy as np
from sklearn.manifold import TSNE
def plot_tsne(model, dataloader, device):
'''
model - torch.nn.Module subclass. This is your encoder model
dataloader - test dataloader to over over data for which you wish to compute projections
device - cuda or cpu (as a string)
'''
model.eval()
images_list = []
labels_list = []
latent_list = []
with torch.no_grad():
for data in dataloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
#approximate the latent space from data
latent_vector = model(images)
images_list.append(images.cpu().numpy())
labels_list.append(labels.cpu().numpy())
latent_list.append(latent_vector.cpu().numpy())
images = np.concatenate(images_list, axis=0)
labels = np.concatenate(labels_list, axis=0)
latent_vectors = np.concatenate(latent_list, axis=0)
# Plot TSNE for latent space
tsne_latent = TSNE(n_components=2, random_state=0)
latent_tsne = tsne_latent.fit_transform(latent_vectors)
plt.figure(figsize=(8, 6))
scatter = plt.scatter(latent_tsne[:, 0], latent_tsne[:, 1], c=labels, cmap='tab10', s=10) # Smaller points
plt.colorbar(scatter)
plt.title('t-SNE of Latent Space')
plt.savefig('latent_tsne.png')
plt.close()
#plot image domain tsne
tsne_image = TSNE(n_components=2, random_state=42)
images_flattened = images.reshape(images.shape[0], -1)
image_tsne = tsne_image.fit_transform(images_flattened)
plt.figure(figsize=(8, 6))
scatter = plt.scatter(image_tsne[:, 0], image_tsne[:, 1], c=labels, cmap='tab10', s=10)
plt.colorbar(scatter)
plt.title('t-SNE of Image Space')
plt.savefig('image_tsne.png')
plt.close() |