A PyTorch implementation of SimCLR based on ICML 2020 paper A Simple Framework for Contrastive Learning of Visual Representations.
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 256 --epochs 500
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 256]
--epochs Number of sweeps over the dataset to train [default value is 500]
python linear.py --batch_size 256 --epochs 200
optional arguments:
--model_path The pretrained model path [default value is 'results/128_0.5_200_256_500_model.pth']
--batch_size Number of images in each mini-batch [default value is 256]
--epochs Number of sweeps over the dataset to train [default value is 200]
There are some difference between this implementation and official implementation, the model (ResNet50
) is trained on
one NVIDIA GeForce RTX 3090(24G) 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.
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 |