SimCLR 对比学习实战TensorFlow Model Garden 中的无监督预训练与半监督微调全流程【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/modelsTensorFlow Model Gardenmode/models仓库中的 SimCLR 项目实现了 SimCLR 与 SimCLRv2 两篇论文提出的对比学习框架先在 ImageNet 上以无监督方式预训练视觉表征再在 1%/10% 标签数据的子集上做半监督微调。读完本文你将掌握两套标准训练命令及其 YAML 配置文件的每个关键参数并能结合official/projects/simclr/下的模型、损失、数据增强源码理解“同一图像两次增强 → 投影头 → 跨副本对比损失”这一核心链路在分布式环境中的具体实现。1. 项目定位与环境准备SimCLRA Simple Framework for Contrastive Learning of Visual Representations, arXiv:2002.05709与 SimCLRv2Big Self-Supervised Models are Strong Semi-Supervised Learners, arXiv:2006.10029的核心思想是对同一张图像施加两次随机增强要求编码器将两个增强视图映射到同一表征方向同时把批内所有其他样本视为负例。official/projects/simclr/README.md是该模块的官方使用文档围绕它我们来看完整的运行方式。1.1 分布式运行环境代码支持多 GPU 或 TPU通过 TensorFlow 的tf.distribute不同分布策略MirroredStrategy、TPU 等运行。文档要求 TensorFlow 2.4全部依赖可一键安装pip install -r ./official/requirements.txt依赖清单见 official/requirements.txt。1.2 训练入口与模式训练统一通过 official/projects/simclr/train.py 启动--mode支持train、eval、train_and_eval等组合。入口的main()流程与 Model Garden 标准训练框架一致gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params)解析 gin 参数train_utils.parse_configuration(FLAGS)根据--experiment名字和--config_fileYAML 构建实验配置对象若mode含train调用train_utils.serialize_config把最终配置序列化写入model_dir纯 eval 模式跳过此步避免与训练任务竞争写文件若设置了params.runtime.mixed_precision_dtype通过performance.set_mixed_precision_policy启用混合精度GPU 上可用float16TPU 上可用bfloat16loss_scale仅在 float16 下生效distribute_utils.get_distribution_strategy(...)按runtime.distribution_strategy、num_gpus、tpu等参数创建分布策略在strategy.scope()内用task_factory.get_task(params.task, ...)构建任务最终交给train_lib.run_experiment执行。此外仓库还提供一个 official/projects/simclr/multitask_train.py 入口与imagenet_simclr_multitask_tpu.yaml配置用于多任务联合训练场景。2. 无监督预训练Pretraining2.1 标准命令在 ImageNet 上预训练模型python3 -m official.projects.simclr.train \ --modetrain_and_eval \ --experimentsimclr_pretraining \ --model_dir{MODEL_DIR} \ --config_file{CONFIG_FILE}其中--experimentsimclr_pretraining对应 configs/simclr.py 中注册的配置工厂exp_factory.register_config_factory(simclr_pretraining) def simclr_pretraining() - cfg.ExperimentConfig: return cfg.ExperimentConfig( taskSimCLRPretrainTask(), trainercfg.TrainerConfig(), restrictions[ task.train_data.is_training ! None, task.validation_data.is_training ! None ])SimCLRPretrainTask的默认值即预训练超参骨架modepretrain、损失为ContrastiveLossprojection_normTrue、temperature0.1支持init_checkpoint/init_checkpoint_modulesall或backbone断点恢复。文件底部还注册了simclr_pretraining_imagenet工厂内置 ResNet-50、batch 4096、500 epoch、LARS 优化器等 ImageNet 完整配方。2.2 预训练配置逐段解读官方 GPU 预训练示例为 configs/experiments/imagenet_simclr_pretrain_gpu.yaml另备有 imagenet_simclr_pretrain_tpu.yaml 与 cifar_simclr_pretrain.yaml。GPU 版配置全文要点runtime 段——多 GPU 混合精度runtime: distribution_strategy: mirrored mixed_precision_dtype: float16 loss_scale: dynamic num_gpus: 16model 段——ResNet-50 主干 投影头 同步 BNtask: model: mode: pretrain input_size: [224, 224, 3] backbone: type: resnet resnet: model_id: 50 backbone_trainable: true projection_head: proj_output_dim: 128 num_proj_layers: 3 ft_proj_idx: 0 supervised_head: num_classes: 1001 norm_activation: use_sync_bn: true norm_momentum: 0.9 norm_epsilon: 0.00001loss / 数据段loss: projection_norm: true temperature: 0.1 train_data: input_path: .../imagenet-2012-tfrecord/train* is_training: true global_batch_size: 2048 dtype: float16 parser: mode: pretrain decoder: decode_label: true示例中的input_path是文档作者的本地 TFRecord 路径使用时替换成自己的imagenet-2012-tfrecord/train*验证集对应valid*。trainer 段——LARS 余弦退火 线性 warmuptrainer: train_steps: 500000 # 800 epochs validation_steps: 24 # NUM_EXAMPLES (50000) // global_batch_size validation_interval: 625 steps_per_loop: 625 # NUM_EXAMPLES (1281167) // global_batch_size summary_interval: 625 checkpoint_interval: 625 optimizer_config: optimizer: type: lars lars: momentum: 0.9 weight_decay_rate: 0.000001 exclude_from_weight_decay: [batch_normalization, bias] learning_rate: type: cosine cosine: initial_learning_rate: 1.6 # 0.2 * BatchSize / 256 decay_steps: 500000 warmup: type: linear linear: warmup_steps: 25000 # 5% of total epochs其中学习率按0.2 × BatchSize / 256线性缩放steps_per_loop 1281167 // 2048 625步即一个 epochtrain_steps: 500000注释表明约为 800 个 epoch 的训练预算。2.3 模型结构为什么一张图要变成两份modeling/simclr_model.py 中SimCLRModel.call揭示了“一份输入、两种输出”的机制。数据管线在预训练模式下把同一图像增强两次后沿通道维拼接为(bsz, h, w, c*2)模型前向时if training and self._mode PRETRAIN: num_transforms 2 # Split channels, and optionally apply extra batched augmentation. # (bsz, h, w, c*num_transforms) - [(bsz, h, w, c), ....] features_list tf.split(inputs, num_or_size_splitsnum_transforms, axis-1) features tf.concat(features_list, 0) # (num_transforms * bsz, h, w, c)随后主干输出接GlobalAveragePooling2D进入投影头同时得到对比学习用的投影向量和微调头用的表征projection_inputs layers.GlobalAveragePooling2D()(features) projection_outputs, supervised_inputs self._projection_head(projection_inputs, training)一个容易忽略的细节预训练阶段即使挂载了supervised_head用于周期性 top-k 评测其输入会做tf.stop_gradient(supervised_inputs)防止监督评测信号污染无监督预训练的梯度。模型还通过checkpoint_items属性把backbone、projection_head、supervised_head分别注册为检查点条目这正是后文init_checkpoint_modules能按模块选择性加载的基础。2.4 投影头中间层有偏置ReLU末层去偏置heads/simclr_head.py 的ProjectionHead.build按num_proj_layers层nn_blocks.DenseBN堆叠中间层use_biasTrue、activationrelu最后一层use_biasFalse、无激活输出维度为proj_output_dim配置中为 128。call返回两个张量hiddens_list[-1]投影头完整输出供对比损失与hiddens_list[self._ft_proj_idx]微调入口特征。ft_proj_idx的语义是0 表示直接取投影头输入骨干池化特征、-1 表示取最后一层输出微调 YAML 中配置为 1即取第一层中间特征。2.5 跨副本对比损失预训练损失在 losses/contrastive_losses.py 的ContrastiveLoss中实现是 SimCLR 大 batch 分布式训练的关键若projection_normTrue先对两路投影做tf.math.l2_normalize多副本时通过cross_replica_concat把各 GPU 上的投影向量all_reduce(SUM)拼成全局 batch。其原理是每份向量先tf.scatter_nd到以replica_id为索引的零张量中再对所有副本求和得到全局张量并在各副本上复制标签构造labels tf.one_hot(tf.range(batch_size), batch_size * 2)指向“自己的另一增强视图”相似度矩阵logits_aa、logits_bb同视图内两两相似度用masks * LARGE_NUMLARGE_NUM 1e9屏蔽自身项logits_ab、logits_ba保留全部跨视图相似度最终损失为两路 softmax 交叉熵之和的均值tf.reduce_mean(loss_a_local loss_b_local)。温度temperature配置中 0.1缩放所有相似度 logits控制负例的“锐度”l2_weight_decay默认 0.0。2.6 数据增强管线预训练与微调的区别dataloaders/simclr_input.py 的Parser._parse_one_train_image定义单条训练图像的增强顺序tf.image.decode_jpeg解码转换为 float [0,1]random_crop_with_resize随机裁剪后 bicubic 缩放回 224×224preprocess_ops.py 中crop_and_resize使用area_range(0.08, 1.0)、aspect_ratio_range(0.75·r, 1.33·r)注释标明 0.875 是 ImageNet 标准裁剪比例tf.image.random_flip_left_right随机水平翻转仅当mode PRETRAINrandom_color_jitter随机颜色扰动color_jitter_strength默认 1.0按0.8×strength调亮度/对比度/饱和度、0.2×strength调色相并以 0.2 概率灰度化仅当mode PRETRAINrandom_blur以 0.5 概率做高斯模糊sigma在 [0.1, 2.0] 间随机、核大小为height // 10bilinear resize、clip 到 [0,1] 并转换为dtype支持 float32/float16/bfloat16。预训练时该函数对同一image_bytes执行两次再tf.concat(xs, -1)得到双视图输入微调时只做一次标准分类预处理。增强开关均可在配置的task.train_data.parser段调整aug_rand_crop、aug_rand_hflip、aug_color_distort、aug_color_jitter_strength、aug_color_jitter_impl可选simclrv1/simclrv2亮度实现、aug_rand_blur、test_crop等默认值与 configs/simclr.py 的Parserdataclass 一致。评估时若test_cropTrue则做中心裁剪输出尺寸 ≤32 时自动关闭适配 CIFAR 类小图数据集。解码器方面Decoder面向 TFRecord解析image/encoded与可选image/class/labelTFDSDecoder面向 TensorFlow Datasets把image重新tf.io.encode_jpeg(quality100)包装后走相同预处理路径这也解释了微调配置可以同时写tfds_name与input_path两个字段。3. 半监督微调1%/10% ImageNet 子集3.1 标准命令README 指出 1% 与 10% ImageNet 子集可通过 TensorFlow Datasets 的imagenet2012_subset数据集直接访问图像 ID 列表位于imagenet_subsets/目录仓库中以tfds_name方式引用。微调整网命令python3 -m official.projects.simclr.train \ --modetrain_and_eval \ --experimentsimclr_finetuning \ --model_dir{MODEL_DIR} \ --config_file{CONFIG_FILE}simclr_finetuning工厂注册SimCLRFinetuneTask默认modefinetune、supervised_headSupervisedHead(num_classes1001, zero_initTrue)、init_checkpoint_modulesbackbone_projection可选all、backbone_projection、backbone损失换为ClassificationLosseslabel_smoothing、one_hot、l2_weight_decay。3.2 微调配置逐段解读官方 10% 微调示例为 configs/experiments/imagenet_simclr_finetune_gpu.yaml另有 imagenet_simclr_finetune_tpu.yaml。与预训练配置相比的关键差异task: model: mode: finetune ... projection_head: proj_output_dim: 128 num_proj_layers: 3 ft_proj_idx: 1 # 从投影头第 1 层取特征接分类头 supervised_head: num_classes: 1001 zero_init: true # 分类头零初始化 norm_activation: use_sync_bn: false # 微调关闭同步 BN loss: label_smoothing: 0.0 one_hot: true init_checkpoint: gs://tf_model_garden/vision/simclr/r50_1x # 预训练检查点 init_checkpoint_modules: backbone_projection train_data: tfds_name: imagenet2012_subset/10pct tfds_split: train global_batch_size: 1024 parser: mode: finetune # 单次增强标准分类预处理trainer: train_steps: 12500 # 100 epochs validation_steps: 49 validation_interval: 125 steps_per_loop: 125 # NUM_EXAMPLES (128116) // global_batch_size optimizer_config: optimizer: type: lars lars: momentum: 0.9 weight_decay_rate: 0.0 # 微调不使用权重衰减 learning_rate: type: cosine cosine: initial_learning_rate: 0.04 # 0.01 × BatchSize / 512 decay_steps: 12500即预训练检查点只加载骨干与投影头backbone_projection分类头零初始化后在 128116 张 10% 子集图像上以 batch 1024 训练 100 个 epoch学习率0.01 × 1024 / 512 0.04余弦退火且无 warmupLARS 权重衰减关闭。改用 1% 子集时把tfds_name换成imagenet2012_subset/1pct即可steps_per_loop等步数需按 12811 张训练样本重新计算。4. 配置体系小结将 README 的命令与仓库实现对照可归纳出一张“参数—源码位置”速查表便于自行裁剪实验配置位置关键参数源码/默认值runtimedistribution_strategy、mixed_precision_dtype、num_gpustrain.py 中distribute_utils/performance.set_mixed_precision_policytask.modelmodepretrain/finetune、backbone、projection_head、supervised_head、backbone_trainablemodeling/simclr_model.py、heads/simclr_head.pytask.losstemperature: 0.1、projection_norm: truelosses/contrastive_losses.pytask.train_data.parser增强开关、aug_color_jitter_impl: simclrv2、test_cropdataloaders/simclr_input.py、dataloaders/preprocess_ops.pytask.init_checkpoint_modulesall/backbone_projection/backboneSimCLRModel.checkpoint_items按模块注册检查点trainerLARS、cosine、linear warmup、steps_per_loopconfigs/simclr.py 的TrainerConfig与优化器工厂对应测试用例 losses/contrastive_losses_test.py、modeling/simclr_model_test.py、heads/simclr_head_test.py、configs/simclr_test.py 覆盖了损失构造、模型双视图拆分与配置工厂的解析行为可作为参数修改后的回归验证手段。5. 复现与引用数据预训练/微调训练集使用imagenet-2012-tfrecordTFRecord 或imagenet2012_subsetDatasets 子集路径通过task.train_data.input_path或tfds_name配置预训练命令用--experimentsimclr_pretraining微调用--experimentsimclr_finetuning两者共用 official/projects/simclr/train.py 入口引用该模块时建议同时引用两篇论文SimCLRChen, Kornblith, Norouzi, Hinton, 2020, arXiv:2002.05709与 SimCLRv2Chen, Kornblith, Swersky, Norouzi, Hinton, 2020, arXiv:2006.10029BibTeX 原文收录在 official/projects/simclr/README.md 的 Cite 小节中。需要提醒的是仓库内的 GPU 配置按 16 块 GPU、batch 2048/1024 设计单卡或少卡复现时应同步下调global_batch_size并按0.2 × BatchSize / 256预训练、0.01 × BatchSize / 512微调规则重算初始学习率同时调整steps_per_loop、train_steps等步数参数以保持与注释中 epoch 数一致。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考