| 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) |
| |
| |
| 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) |
| |
| |
| 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) |
| plt.colorbar(scatter) |
| plt.title('t-SNE of Latent Space') |
| plt.savefig('latent_tsne.png') |
| plt.close() |
| |
| |
| 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() |
