mirrors/SimCLR
1
0
Fork 0
mirror of https://github.com/sthalles/SimCLR.git synced 2025-06-03 15:03:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Major refactor, small fixes

Browse Source
This commit is contained in:
Thalles Silva 2021-01-18 07:33:12 -03:00
parent f78dd0a792
commit 6cd188ef68
3 changed files with 19 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
README.md
View File

@ -16,78 +16,27 @@ $ python run.py
## Config file ## Config file
Before running SimCLR, make sure you choose the correct running configurations on the ```config.yaml``` file. Before running SimCLR, make sure you choose the correct running configurations. You can change the running configurations by passing keyword arguments to the ```run.py``` file.
```yaml ```python
# A batch size of N, produces 2 * (N-1) negative samples. Original implementation uses a batch size of 8192 $ python run.py -data ./datasets --dataset-name stl10 --log-every-n-steps 100 --epochs 100
batch_size: 512
# Number of epochs to train
epochs: 40
# Frequency to eval the similarity score using the validation set
eval_every_n_epochs: 1
# Specify a folder containing a pre-trained model to fine-tune. If training from scratch, pass None.
fine_tune_from: 'resnet-18_80-epochs'
# Frequency to which tensorboard is updated
log_every_n_steps: 50
# l2 Weight decay magnitude, original implementation uses 10e-6
weight_decay: 10e-6
# if True, training is done using mixed precision. Apex needs to be installed in this case.
fp16_precision: False
# Model related parameters
model:
# Output dimensionality of the embedding vector z. Original implementation uses 2048
out_dim: 256
# The ConvNet base model. Choose one of: "resnet18" or "resnet50". Original implementation uses resnet50
base_model: "resnet18"
# Dataset related parameters
dataset:
s: 1
# dataset input shape. For datasets containing images of different size, this defines the final
input_shape: (96,96,3)
# Number of workers for the data loader
num_workers: 0
# Size of the validation set in percentage
valid_size: 0.05
# NTXent loss related parameters
loss:
# Temperature parameter for the contrastive objective
temperature: 0.5
# Distance metric for contrastive loss. If False, uses dot product. Original implementation uses cosine similarity.
use_cosine_similarity: True
``` ```
If you want to run it on CPU (for debugging purposes) use the ```--disable-cuda``` option.
## Feature Evaluation ## Feature Evaluation
Feature evaluation is done using a linear model protocol. Feature evaluation is done using a linear model protocol.
Features are learned using the ```STL10 train+unsupervised``` set and evaluated in the ```test``` set; First, we learned features using SimCLR on the ```STL10 unsupervised``` set. Then, we train a linear classifier on top of the frozen features from SimCLR. The linera model is trained on features extracted from the ```STL10 train``` set and evaluated on the ```STL10 test``` set.
Check the [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/sthalles/SimCLR/blob/master/feature_eval/linear_feature_eval.ipynb) notebook for reproducibility. Check the [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/sthalles/SimCLR/blob/simclr-refactor/feature_eval/mini_batch_logistic_regression_evaluator.ipynb) notebook for reproducibility.
| Linear Classifier | Feature Extractor | Architecture | Feature dimensionality | Projection Head dimensionality | Epochs | STL10 Top 1 | | Linear Classification | Dataset | Feature Extractor | Architecture | Feature dimensionality | Projection Head dimensionality | Epochs | Top 1 |
|:---------------------------:|:-----------------:|:------------:|:----------------------:|:-------------------------------:|:------:|:-----------:| |----------------------------|---------|-------------------|---------------------------------------------------------------------------------|------------------------|--------------------------------|--------|--------|
| Logistic Regression | PCA Features | - | 256 | - | | 36.0% | | Logistic Regression (Adam) | STL10 | SimCLR | [ResNet-18](https://drive.google.com/open?id=14_nH2FkyKbt61cieQDiSbBVNP8-gtwgF) | 512 | 128 | 100 | 70.45 |
| KNN | PCA Features | - | 256 | - | | 31.8% | | Logistic Regression (Adam) | CIFAR10 | SimCLR | [ResNet-18](https://drive.google.com/open?id=1lc2aoVtrAetGn0PnTkOyFzPCIucOJq7C) | 512 | 128 | 100 | 64.82 |
| Logistic Regression (LBFGS) | SimCLR | [ResNet-18](https://drive.google.com/open?id=1c4eVon0sUd-ChVhH6XMpF6nCngNJsAPk) | 512 | 256 | 40 | 70.3% | | Logistic Regression (Adam) | STL10 | SimCLR | [ResNet-50](https://drive.google.com/open?id=1ByTKAUsdm_X7tLcii6oAEl5qFRqRMZSu) | 2048 | 128 | 50 | 67.075 |
| KNN | SimCLR | ResNet-18 | 512 | 256 | 40 | 66.2% |
| Logistic Regression (LBFGS) | SimCLR | [ResNet-18](https://drive.google.com/open?id=1L0yoeY9i2mzDcj69P4slTWb-cfr3PyoT) | 512 | 256 | 80 | 72.9% |
| KNN | SimCLR | ResNet-18 | 512 | 256 | 80 | 69.8% |
| Logistic Regression (Adam) | SimCLR | [ResNet-18](https://drive.google.com/open?id=1aZ12TITXnajZ6QWmS_SDm8Sp8gXNbeCQ) | 512 | 256 | 100 | 75.4% |
| Logistic Regression (Adam) | SimCLR | [ResNet-50](https://drive.google.com/open?id=1TZqBNTFCsO-mxAiR-zJeyupY-J2gA27Q) | 2048 | 128 | 40 | 74.6% |
| Logistic Regression (Adam) | SimCLR | [ResNet-50](https://drive.google.com/open?id=1is1wkBRccHdhSKQnPUTQoaFkVNSaCb35) | 2048 | 128 | 80 | 77.3% |

4
run.py
View File

@ -1,10 +1,10 @@
import argparse import argparse
import torch import torch
import torch.backends.cudnn as cudnn
from torchvision import models from torchvision import models
from data_aug.contrastive_learning_dataset import ContrastiveLearningDataset from data_aug.contrastive_learning_dataset import ContrastiveLearningDataset
from models.resnet_simclr import ResNetSimCLR from models.resnet_simclr import ResNetSimCLR
from simclr import SimCLR from simclr import SimCLR
import torch.backends.cudnn as cudnn
model_names = sorted(name for name in models.__dict__ model_names = sorted(name for name in models.__dict__
if name.islower() and not name.startswith("__") if name.islower() and not name.startswith("__")
@ -34,8 +34,6 @@ parser.add_argument('--lr', '--learning-rate', default=0.0003, type=float,
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float, parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)', metavar='W', help='weight decay (default: 1e-4)',
dest='weight_decay') dest='weight_decay')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('--seed', default=None, type=int, parser.add_argument('--seed', default=None, type=int,
help='seed for initializing training. ') help='seed for initializing training. ')
parser.add_argument('--disable-cuda', action='store_true', parser.add_argument('--disable-cuda', action='store_true',

12
simclr.py
View File

@ -1,11 +1,12 @@
import logging
import os import os
import shutil import shutil
import sys import sys
import yaml
import torch import torch
from torch.utils.tensorboard import SummaryWriter
import torch.nn.functional as F import torch.nn.functional as F
import logging import yaml
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm from tqdm import tqdm
torch.manual_seed(0) torch.manual_seed(0)
@ -84,11 +85,8 @@ class SimCLR(object):
self.model, self.optimizer = amp.initialize(self.model, self.optimizer, self.model, self.optimizer = amp.initialize(self.model, self.optimizer,
opt_level='O2', opt_level='O2',
keep_batchnorm_fp32=True) keep_batchnorm_fp32=True)
model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints')
# save config file # save config file
_save_config_file(model_checkpoints_folder, self.args) _save_config_file(self.writer.log_dir, self.args)
n_iter = 0 n_iter = 0
logging.info(f"Start SimCLR training for {self.args.epochs} epochs.") logging.info(f"Start SimCLR training for {self.args.epochs} epochs.")