Initial Project

code/environment.yml

@@ -0,0 +1,51 @@
+name: aes
+channels:
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - ca-certificates=2024.12.31=h06a4308_0
+  - ld_impl_linux-64=2.40=h12ee557_0
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - ncurses=6.4=h6a678d5_0
+  - openssl=1.1.1w=h7f8727e_0
+  - pip=25.0=py39h06a4308_0
+  - python=3.9.0=hdb3f193_2
+  - readline=8.2=h5eee18b_0
+  - setuptools=75.8.0=py39h06a4308_0
+  - sqlite=3.45.3=h5eee18b_0
+  - tk=8.6.14=h39e8969_0
+  - tzdata=2025a=h04d1e81_0
+  - wheel=0.45.1=py39h06a4308_0
+  - xz=5.4.6=h5eee18b_1
+  - zlib=1.2.13=h5eee18b_1
+  - pip:
+      - contourpy==1.3.0
+      - cycler==0.12.1
+      - filelock==3.17.0
+      - fonttools==4.55.8
+      - fsspec==2025.2.0
+      - importlib-resources==6.5.2
+      - jinja2==3.1.5
+      - kiwisolver==1.4.7
+      - markupsafe==3.0.2
+      - matplotlib==3.9.4
+      - mpmath==1.3.0
+      - networkx==3.2.1
+      - numpy==2.0.2
+      - packaging==24.2
+      - pillow==11.1.0
+      - pyparsing==3.2.1
+      - python-dateutil==2.9.0.post0
+      - six==1.17.0
+      - sympy==1.13.1
+      - torch==2.6.0
+      - torchaudio==2.6.0
+      - torchvision==0.21.0
+      - triton==3.2.0
+      - typing-extensions==4.12.2
+      - zipp==3.21.0
+      - scikit-learn

code/main.py

@@ -0,0 +1,56 @@
+import torch
+from torchvision import datasets, transforms
+import numpy as np
+from matplotlib import pyplot as plt
+from utils import plot_tsne
+import numpy as np
+import random
+import argparse
+
+NUM_CLASSES = 10
+
+def freeze_seeds(seed=0):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+def get_args():   
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--seed', default=0, type=int, help='Seed for random number generators')
+    parser.add_argument('--data-path', default="/datasets/cv_datasets/data", type=str, help='Path to dataset')
+    parser.add_argument('--batch-size', default=8, type=int, help='Size of each batch')
+    parser.add_argument('--latent-dim', default=128, type=int, help='encoding dimension')
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help='Default device to use')
+    parser.add_argument('--mnist', action='store_true', default=False,
+                        help='Whether to use MNIST (True) or CIFAR10 (False) data')
+    parser.add_argument('--self-supervised', action='store_true', default=False,
+                        help='Whether train self-supervised with reconstruction objective, or jointly with classifier for classification objective.')
+    return parser.parse_args()
+    
+
+if __name__ == "__main__":
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  #one possible convenient normalization. You don't have to use it.
+    ])
+
+    args = get_args()
+    freeze_seeds(args.seed)
+                
+                                           
+    if args.mnist:
+        train_dataset = datasets.MNIST(root=args.data_path, train=True, download=False, transform=transform)
+        test_dataset = datasets.MNIST(root=args.data_path, train=False, download=False, transform=transform)
+    else:
+        train_dataset = datasets.CIFAR10(root=args.data_path, train=True, download=True, transform=transform)
+        test_dataset = datasets.CIFAR10(root=args.data_path, train=False, download=True, transform=transform)
+        
+    # When you create your dataloader you should split train_dataset or test_dataset to leave some aside for validation
+
+    #this is just for the example. Simple flattening of the image is probably not the best idea                                        
+    encoder_model = torch.nn.Linear(32*32*3,args.latent_dim).to(args.device)
+    decoder_model = torch.nn.Linear(args.latent_dim,32*32*3 if args.self_supervised else NUM_CLASSES).to(args.device) 
+
+    sample = train_dataset[0][0][None].to(args.device) #This is just for the example - you should use a dataloader
+    output = decoder_model(encoder_model(sample.flatten()))
+    print(output.shape)
+

code/utils.py

@@ -0,0 +1,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()