Notícias! Lançamos uma implementação TF2 do SimCLR (junto com pontos de verificação convertidos no TF2), eles estão na pasta tf2/.
Notícias! Foram adicionados Colabs para propriedades intrigantes de perdas contrastivas, veja aqui.
Abrimos um total de 65 modelos pré-treinados aqui, correspondentes aos da Tabela 1 do artigo SimCLRv2:
Profundidade | Largura | SK | Parâmetro (M) | TF (1%) | Pés (10%) | Pés (100%) | Avaliação linear | Supervisionado |
---|---|---|---|---|---|---|---|---|
50 | 1X | Falso | 24 | 57,9 | 68,4 | 76,3 | 71,7 | 76,6 |
50 | 1X | Verdadeiro | 35 | 64,5 | 72,1 | 78,7 | 74,6 | 78,5 |
50 | 2X | Falso | 94 | 66,3 | 73,9 | 79,1 | 75,6 | 77,8 |
50 | 2X | Verdadeiro | 140 | 70,6 | 77,0 | 81,3 | 77,7 | 79,3 |
101 | 1X | Falso | 43 | 62,1 | 71,4 | 78,2 | 73,6 | 78,0 |
101 | 1X | Verdadeiro | 65 | 68,3 | 75,1 | 80,6 | 76,3 | 79,6 |
101 | 2X | Falso | 170 | 69,1 | 75,8 | 80,7 | 77,0 | 78,9 |
101 | 2X | Verdadeiro | 257 | 73,2 | 78,8 | 82,4 | 79,0 | 80,1 |
152 | 1X | Falso | 58 | 64,0 | 73,0 | 79,3 | 74,5 | 78,3 |
152 | 1X | Verdadeiro | 89 | 70,0 | 76,5 | 81,3 | 77,2 | 79,9 |
152 | 2X | Falso | 233 | 70,2 | 76,6 | 81,1 | 77,4 | 79,1 |
152 | 2X | Verdadeiro | 354 | 74,2 | 79,4 | 82,9 | 79,4 | 80,4 |
152 | 3X | Verdadeiro | 795 | 74,9 | 80,1 | 83,1 | 79,8 | 80,5 |
Esses pontos de verificação são armazenados no Google Cloud Storage:
Também fornecemos exemplos de como usar os pontos de verificação na pasta colabs/
.
Os modelos pré-treinados (rede base com camada classificadora linear) podem ser encontrados abaixo. Observe que para esses pontos de verificação SimCLRv1, o cabeçote de projeção não está disponível.
Ponto de verificação do modelo e módulo hub | ImageNet Top 1 |
---|---|
ResNet50 (1x) | 69,1 |
ResNet50 (2x) | 74,2 |
ResNet50 (4x) | 76,6 |
Pontos de verificação SimCLRv1 adicionais estão disponíveis: gs://simclr-checkpoints/simclrv1.
Uma observação sobre as assinaturas do módulo TensorFlow Hub: default
é a saída de representação da rede base; logits_sup
são os logits de classificação supervisionados para categorias do ImageNet 1000. Outros (por exemplo, initial_max_pool
, block_group1
) são camadas intermediárias do ResNet; consulte resnet.py para obter detalhes. Consulte este tutorial para obter informações adicionais sobre o uso dos módulos do TensorFlow Hub.
Nossos modelos são treinados com TPUs. Recomenda-se executar treinamento distribuído com TPUs ao usar nosso código para pré-treinamento.
Nosso código também pode ser executado em uma única GPU. Ele não oferece suporte a multi-GPUs, por motivos como BatchNorm global e perda contrastiva entre núcleos.
O código é compatível com TensorFlow v1 e v2. Consulte requisitos.txt para todos os pré-requisitos e você também pode instalá-los usando o comando a seguir.
pip install -r requirements.txt
Para pré-treinar o modelo no CIFAR-10 com uma única GPU, tente o seguinte comando:
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
Para pré-treinar o modelo no ImageNet com Cloud TPUs, primeiro confira o tutorial do Google Cloud TPU para obter informações básicas sobre como usar o Google Cloud TPUs.
Depois de criar a máquina virtual com Cloud TPUs e fazer o pré-download dos dados do ImageNet para tensorflow_datasets, defina as seguintes variáveis de ambiente:
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>
O comando a seguir pode ser usado para pré-treinar um ResNet-50 no ImageNet (que reflete os hiperparâmetros padrão em nosso artigo):
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
Um tamanho de lote de 4.096 requer pelo menos 32 TPUs. 100 épocas levam cerca de 6 horas com 32 TPU v3s. Observe que a taxa de aprendizado de 0,3 com learning_rate_scaling=linear
é equivalente a 0,075 com learning_rate_scaling=sqrt
quando o tamanho do lote é 4096. No entanto, usar o escalonamento sqrt permite treinar melhor quando um tamanho de lote menor é usado.
Para ajustar uma cabeça linear (com uma única GPU), tente o seguinte comando:
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
Você pode verificar os resultados usando tensorboard, como
python -m tensorboard.main --logdir=/tmp/simclr_test
Como referência, as execuções acima no CIFAR-10 devem fornecer cerca de 91% de precisão, embora possam ser otimizadas ainda mais.
Para ajustar um cabeçote linear no ImageNet usando Cloud TPUs, primeiro defina CHKPT_DIR
como diretório do modelo pré-treinado e defina um novo MODEL_DIR
e, em seguida, use o seguinte comando:
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
Como referência, as execuções acima no ImageNet devem fornecer cerca de 64,5% de precisão.
Você pode acessar subconjuntos ImageNet de 1% e 10% usados para aprendizado semissupervisionado por meio de conjuntos de dados tensorflow: basta definir dataset=imagenet2012_subset/1pct
e dataset=imagenet2012_subset/10pct
na linha de comando para ajuste fino nesses subconjuntos.
Você também pode encontrar IDs de imagem desses subconjuntos em imagenet_subsets/
.
Para ajustar toda a rede no ImageNet (1% dos rótulos), consulte o seguinte comando:
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
Defina o checkpoint
para aqueles que são apenas pré-treinados, mas não ajustados. Dado que os pontos de verificação SimCLRv1 não contêm cabeça de projeção, é recomendado executar com pontos de verificação SimCLRv2 (você ainda pode executar com pontos de verificação SimCLRv1, mas variable_schema
precisa excluir head
). O num_proj_layers
e ft_proj_selector
precisam ser ajustados de acordo com o artigo SimCLRv2 para obter melhores desempenhos.
Este repositório fornece uma solução para converter os pontos de verificação SimCLRv1 Tensorflow pré-treinados em pontos Pytorch.
Este repositório fornece uma solução para converter os pontos de verificação SimCLRv2 Tensorflow pré-treinados em pontos Pytorch.
(Sinta-se à vontade para compartilhar sua implementação criando um problema)
Implementações em PyTorch:
Implementações no Tensorflow 2/Keras (a implementação oficial do TF2 foi adicionada na pasta tf2/):
Tamanho do lote : os resultados originais do SimCLR foram ajustados para um tamanho de lote grande (ou seja, 4096), o que leva a resultados abaixo do ideal ao treinar usando um tamanho de lote menor. No entanto, com um bom conjunto de hiperparâmetros (principalmente taxa de aprendizagem, temperatura, profundidade da cabeça de projeção), lotes pequenos podem produzir resultados equivalentes a lotes grandes (por exemplo, consulte a Tabela 2 neste artigo).
Modelos/pontos de verificação pré-treinados : SimCLRv1 e SimCLRv2 são pré-treinados com diferentes reduções de peso, portanto, os modelos pré-treinados das duas versões têm escalas de norma de peso muito diferentes (os pesos convolucionais no SimCLRv1 ResNet-50 são em média 16,8X daqueles no SimCLRv2). Para ajustar os modelos pré-treinados de ambas as versões, não há problema se você usar um otimizador LARS, mas serão necessários hiperparâmetros muito diferentes (por exemplo, taxa de aprendizagem, redução de peso) se você usar o otimizador de momento. Portanto, para o último caso, você pode querer procurar hparams muito diferentes de acordo com a versão usada ou redimensionar o peso (ou seja, conv parâmetros kernel
de base_model
nos pontos de verificação) para ter certeza de que estão aproximadamente na mesma escala .
Artigo 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}
}
Artigo 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}
}
Este não é um produto oficial do Google.