Awesome
SimCLR
A PyTorch implementation of SimCLR based on ICML 2020 paper A Simple Framework for Contrastive Learning of Visual Representations.
Requirements
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
- thop
pip install thop
Dataset
CIFAR10
dataset is used in this repo, the dataset will be downloaded into data
directory by PyTorch
automatically.
Usage
Train SimCLR
python main.py --batch_size 1024 --epochs 1000
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]
Linear Evaluation
python linear.py --batch_size 1024 --epochs 200
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]
Results
There are some difference between this implementation and official implementation, the model (ResNet50
) is trained on
one NVIDIA TESLA V100(32G) GPU:
- No
Gaussian blur
used; Adam
optimizer with learning rate1e-3
is used to replaceLARS
optimizer;- No
Linear learning rate scaling
used; - No
Linear Warmup
andCosineLR Schedule
used.