A PyTorch implementation of SimCLR based on ICML 2020 paper A Simple Framework for Contrastive Learning of Visual Representations.
Network Architecture image from the paper
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
- thop
pip install thop
CIFAR10 dataset is used in this repo, the dataset will be downloaded into data directory by PyTorch automatically.
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]
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]
There are some difference between this implementation and official implementation, the model (ResNet50) is trained on
one NVIDIA TESLA V100(32G) GPU:
- No
Gaussian blurused; Adamoptimizer with learning rate1e-3is used to replaceLARSoptimizer;- No
Linear learning rate scalingused; - No
Linear WarmupandCosineLR Scheduleused.
| Evaluation Protocol | Params (M) | FLOPs (G) | Feature Dim | Batch Size | Epoch Num | τ | K | Top1 Acc % | Top5 Acc % | Download |
|---|---|---|---|---|---|---|---|---|---|---|
| KNN | 24.62 | 1.31 | 128 | 512 | 500 | 0.5 | 200 | 89.1 | 99.6 | model | gc5k |
| Linear | 23.52 | 1.30 | - | 512 | 100 | - | - | 92.0 | 99.8 | model | f7j2 |