Cooper Ang
Accompanied paper is availiable as a PDF here.
This repository contains the code for the final project of SYDE 572. The goal of this project was to implement a self-supervised contrastive learning algorithm and apply it with smaller more efficient model architectures. The code is based on the SimCLR paper.
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
- thop
- matplotlib
- numpy
- pandas
pip install thop
pip install matplotlib
pip install numpy
pip install pandas
CIFAR10 dataset is used in this repo, the dataset will be downloaded into data directory by PyTorch automatically.
python main.py
optional arguments:
--feature_dim Feature dim for latent vector [default value is 128]
--temperature Temperature used in softmax [default value is 0.5]
--k Top k most similar images used to predict the label [default value is 200]
--batch_size Number of images in each mini-batch [default value is 512]
--epochs Number of sweeps over the dataset to train [default value is 500]
--model_base The base model used for training [default value is 'mobilenetv3_small'] Options are 'mobilenetv3_small', 'mobilenetv3_large', 'resnet18'
python linear.py
optional arguments:
--model_path The pretrained model path [default value is 'results/128_0.5_200_512_500_model.pth']
--batch_size Number of images in each mini-batch [default value is 512]
--epochs Number of sweeps over the dataset to train [default value is 100]
--model_base The base modle used for training [default value is 'mobilenetv3_small'] Options are 'mobilenetv3_small', 'mobilenetv3_large', 'resnet18'
Details are in the `Generate Plots.ipynb` notebook
Detailed results can be found in final_results directory.
Nested inside the directory corresponds to different model candidates with self-explanatory titles corresponding to the base model architecture and the hyperparameters used in the training.
This repo is based on the following repo which implemented the SimCLR algorithm: