ニュース! SimCLR の TF2 実装を (TF2 で変換されたチェックポイントとともに) リリースしました。それらは tf2/ フォルダーにあります。
ニュース!対照的な損失の興味深い特性に関するコラボが追加されました。こちらを参照してください。
ここでは、SimCLRv2 論文の表 1 にあるモデルに対応する、合計 65 個の事前トレーニング済みモデルをオープンソース化しました。
深さ | 幅 | SK | パラム(男) | FT (1%) | FT(10%) | FT(100%) | 線形評価 | 監督あり |
---|---|---|---|---|---|---|---|---|
50 | 1X | 間違い | 24 | 57.9 | 68.4 | 76.3 | 71.7 | 76.6 |
50 | 1X | 真実 | 35 | 64.5 | 72.1 | 78.7 | 74.6 | 78.5 |
50 | 2倍 | 間違い | 94 | 66.3 | 73.9 | 79.1 | 75.6 | 77.8 |
50 | 2倍 | 真実 | 140 | 70.6 | 77.0 | 81.3 | 77.7 | 79.3 |
101 | 1X | 間違い | 43 | 62.1 | 71.4 | 78.2 | 73.6 | 78.0 |
101 | 1X | 真実 | 65 | 68.3 | 75.1 | 80.6 | 76.3 | 79.6 |
101 | 2倍 | 間違い | 170 | 69.1 | 75.8 | 80.7 | 77.0 | 78.9 |
101 | 2倍 | 真実 | 257 | 73.2 | 78.8 | 82.4 | 79.0 | 80.1 |
152 | 1X | 間違い | 58 | 64.0 | 73.0 | 79.3 | 74.5 | 78.3 |
152 | 1X | 真実 | 89 | 70.0 | 76.5 | 81.3 | 77.2 | 79.9 |
152 | 2倍 | 間違い | 233 | 70.2 | 76.6 | 81.1 | 77.4 | 79.1 |
152 | 2倍 | 真実 | 354 | 74.2 | 79.4 | 82.9 | 79.4 | 80.4 |
152 | 3X | 真実 | 795 | 74.9 | 80.1 | 83.1 | 79.8 | 80.5 |
これらのチェックポイントは Google Cloud Storage に保存されます。
colabs/
フォルダー内のチェックポイントの使用方法の例も提供します。
事前トレーニング済みモデル (線形分類層を備えたベース ネットワーク) は以下にあります。これらの SimCLRv1 チェックポイントでは、プロジェクション ヘッドが利用できないことに注意してください。
モデルのチェックポイントとハブモジュール | イメージネットトップ1 |
---|---|
ResNet50 (1x) | 69.1 |
ResNet50 (2x) | 74.2 |
ResNet50 (4x) | 76.6 |
追加の SimCLRv1 チェックポイントが利用可能です: gs://simclr-checkpoints/simclrv1。
TensorFlow Hub モジュールの署名に関するメモ: default
ベース ネットワークの表現出力です。 logits_sup
、ImageNet 1000 カテゴリの教師付き分類ロジットです。その他 (例: initial_max_pool
、 block_group1
)は ResNet の中間層です。詳細については、resnet.py を参照してください。 TensorFlow Hub モジュールの使用に関する追加情報については、このチュートリアルを参照してください。
私たちのモデルは TPU でトレーニングされています。事前トレーニングにコードを使用する場合は、TPU を使用して分散トレーニングを実行することをお勧めします。
私たちのコードは単一のGPU 上でも実行できます。グローバル BatchNorm やコア間の対照的な損失などの理由により、マルチ GPU はサポートされていません。
このコードは TensorFlow v1 と v2 の両方と互換性があります。すべての前提条件については、requirements.txt を参照してください。また、次のコマンドを使用してインストールすることもできます。
pip install -r requirements.txt
単一のGPU を使用して CIFAR-10 でモデルを事前トレーニングするには、次のコマンドを試してください。
python run.py --train_mode=pretrain
--train_batch_size=512 --train_epochs=1000
--learning_rate=1.0 --weight_decay=1e-4 --temperature=0.5
--dataset=cifar10 --image_size=32 --eval_split=test --resnet_depth=18
--use_blur=False --color_jitter_strength=0.5
--model_dir=/tmp/simclr_test --use_tpu=False
Cloud TPU を使用して ImageNet でモデルを事前トレーニングするには、まず Google Cloud TPU チュートリアルを参照して、Google Cloud TPU の使用方法に関する基本情報を確認してください。
Cloud TPU を使用して仮想マシンを作成し、tensorflow_datasets の ImageNet データを事前にダウンロードしたら、次の環境変数を設定してください。
TPU_NAME=<tpu-name>
STORAGE_BUCKET=gs://<storage-bucket>
DATA_DIR=$STORAGE_BUCKET/<path-to-tensorflow-dataset>
MODEL_DIR=$STORAGE_BUCKET/<path-to-store-checkpoints>
次のコマンドを使用して、ImageNet で ResNet-50 を事前トレーニングできます (論文のデフォルトのハイパーパラメーターを反映しています)。
python run.py --train_mode=pretrain
--train_batch_size=4096 --train_epochs=100 --temperature=0.1
--learning_rate=0.075 --learning_rate_scaling=sqrt --weight_decay=1e-4
--dataset=imagenet2012 --image_size=224 --eval_split=validation
--data_dir=$DATA_DIR --model_dir=$MODEL_DIR
--use_tpu=True --tpu_name=$TPU_NAME --train_summary_steps=0
バッチ サイズ 4096 には、少なくとも 32 TPU が必要です。 100 エポックには 32 TPU v3 で約 6 時間かかります。 learning_rate_scaling=linear
の学習率 0.3 は、バッチ サイズが 4096 の場合、 learning_rate_scaling=sqrt
の 0.075 と同等であることに注意してください。ただし、sqrt スケーリングを使用すると、より小さいバッチ サイズが使用される場合に、より適切にトレーニングできます。
リニア ヘッド (単一 GPU を使用) を微調整するには、次のコマンドを試してください。
python run.py --mode=train_then_eval --train_mode=finetune
--fine_tune_after_block=4 --zero_init_logits_layer=True
--variable_schema='(?!global_step|(?:.*/|^)Momentum|head)'
--global_bn=False --optimizer=momentum --learning_rate=0.1 --weight_decay=0.0
--train_epochs=100 --train_batch_size=512 --warmup_epochs=0
--dataset=cifar10 --image_size=32 --eval_split=test --resnet_depth=18
--checkpoint=/tmp/simclr_test --model_dir=/tmp/simclr_test_ft --use_tpu=False
次のようなテンソルボードを使用して結果を確認できます。
python -m tensorboard.main --logdir=/tmp/simclr_test
参考として、CIFAR-10 で上記を実行すると、約 91% の精度が得られますが、さらに最適化することもできます。
Cloud TPU を使用して ImageNet 上のリニア ヘッドを微調整するには、まずCHKPT_DIR
事前トレーニングされたモデル ディレクトリに設定し、新しいMODEL_DIR
を設定してから、次のコマンドを使用します。
python run.py --mode=train_then_eval --train_mode=finetune
--fine_tune_after_block=4 --zero_init_logits_layer=True
--variable_schema='(?!global_step|(?:.*/|^)Momentum|head)'
--global_bn=False --optimizer=momentum --learning_rate=0.1 --weight_decay=1e-6
--train_epochs=90 --train_batch_size=4096 --warmup_epochs=0
--dataset=imagenet2012 --image_size=224 --eval_split=validation
--data_dir=$DATA_DIR --model_dir=$MODEL_DIR --checkpoint=$CHKPT_DIR
--use_tpu=True --tpu_name=$TPU_NAME --train_summary_steps=0
参考として、ImageNet で上記を実行すると、約 64.5% の精度が得られます。
tensorflow データセットを介して、半教師あり学習に使用される 1% および 10% の ImageNet サブセットにアクセスできます。これらのサブセットを微調整するには、コマンド ラインでdataset=imagenet2012_subset/1pct
およびdataset=imagenet2012_subset/10pct
を設定するだけです。
これらのサブセットのイメージ ID は、 imagenet_subsets/
で見つけることもできます。
ImageNet 上のネットワーク全体 (ラベルの 1%) を微調整するには、次のコマンドを参照してください。
python run.py --mode=train_then_eval --train_mode=finetune
--fine_tune_after_block=-1 --zero_init_logits_layer=True
--variable_schema='(?!global_step|(?:.*/|^)Momentum|head_supervised)'
--global_bn=True --optimizer=lars --learning_rate=0.005
--learning_rate_scaling=sqrt --weight_decay=0
--train_epochs=60 --train_batch_size=1024 --warmup_epochs=0
--dataset=imagenet2012_subset/1pct --image_size=224 --eval_split=validation
--data_dir=$DATA_DIR --model_dir=$MODEL_DIR --checkpoint=$CHKPT_DIR
--use_tpu=True --tpu_name=$TPU_NAME --train_summary_steps=0
--num_proj_layers=3 --ft_proj_selector=1
checkpoint
、事前トレーニングのみで微調整されていないものに設定します。 SimCLRv1 チェックポイントには射影 head が含まれていないため、SimCLRv2 チェックポイントを使用して実行することをお勧めします (SimCLRv1 チェックポイントを使用しても実行できますが、 variable_schema
head
除外する必要があります)。最高のパフォーマンスを得るには、SimCLRv2 の論文に従ってnum_proj_layers
とft_proj_selector
調整する必要があります。
このリポジトリは、事前トレーニング済みの SimCLRv1 Tensorflow チェックポイントを Pytorch チェックポイントに変換するためのソリューションを提供します。
このリポジトリは、事前トレーニング済みの SimCLRv2 Tensorflow チェックポイントを Pytorch チェックポイントに変換するためのソリューションを提供します。
(問題を作成して実装を自由に共有してください)
PyTorch での実装:
Tensorflow 2 / Keras での実装 (公式の TF2 実装は tf2/ フォルダーに追加されました):
バッチ サイズ: SimCLR の元の結果は大きなバッチ サイズ (つまり 4096) で調整されているため、より小さいバッチ サイズを使用してトレーニングすると次善の結果が得られます。ただし、ハイパーパラメータの適切なセット (主に学習率、温度、投影ヘッドの深さ) を使用すると、小さなバッチ サイズでも大きなバッチ サイズと同等の結果が得られます (たとえば、この論文の表 2 を参照)。
事前トレーニング済みモデル / チェックポイント: SimCLRv1 と SimCLRv2 は異なる重み減衰で事前トレーニングされているため、2 つのバージョンの事前トレーニング済みモデルの重みノルム スケールは大きく異なります (SimCLRv1 ResNet-50 の畳み込み重みは、SimCLRv2 の重みの平均 16.8 倍です)。両方のバージョンの事前トレーニング済みモデルを微調整する場合、LARS オプティマイザーを使用する場合は問題ありませんが、モメンタム オプティマイザーを使用する場合は、非常に異なるハイパーパラメーター (学習率、重み減衰など) が必要になります。したがって、後者の場合は、使用されているバージョンに応じて非常に異なる hparams を検索するか、重み (つまり、チェックポイントのbase_model
の変換kernel
パラメーター) を再スケールして、それらがほぼ同じスケールであることを確認する必要があるかもしれません。 。
SimCLR 論文:
@article{chen2020simple,
title={A Simple Framework for Contrastive Learning of Visual Representations},
author={Chen, Ting and Kornblith, Simon and Norouzi, Mohammad and Hinton, Geoffrey},
journal={arXiv preprint arXiv:2002.05709},
year={2020}
}
SimCLRv2 論文:
@article{chen2020big,
title={Big Self-Supervised Models are Strong Semi-Supervised Learners},
author={Chen, Ting and Kornblith, Simon and Swersky, Kevin and Norouzi, Mohammad and Hinton, Geoffrey},
journal={arXiv preprint arXiv:2006.10029},
year={2020}
}
これは Google の公式製品ではありません。