冷酷的教师是更好的榜样

论文 ---- NIPS 2017 海报 ---- NIPS 2017 焦点演讲幻灯片 ---- 博客文章

作者：Antti Tarvainen、Harri Valpola（好奇人工智能公司）

方法

均值教师是一种简单的半监督学习方法。其步骤如下：

1. 拿一个有监督的模型架构并复制一份。我们称原始模型为“学生”，新模型为“教师”。
2. 在每个训练步骤中，将相同的迷你批次分别输入到学生和教师模型中，但对输入分别添加随机增强或噪声。
3. 在学生和教师的输出之间（经过 softmax 后）加入额外的一致性损失项。
4. 让优化器正常更新学生的权重。
5. 教师的权重设置为学生权重的指数移动平均（EMA）。也就是说，在每个训练步骤后，教师的权重会稍微向学生权重靠拢。

我们的贡献在于最后一步。Laine 和 Aila [论文] 使用了学生和教师之间的共享参数，或者使用了教师预测的时间集成。相比之下，均值教师更加准确，并且适用于大规模数据集。

均值教师与现代架构配合良好。通过将均值教师与 ResNet 结合，我们在 ImageNet 和 CIFAR-10 数据集上的半监督学习性能达到了新的高度。

ImageNet 使用 10% 标签	验证集 top-5 错误率
变分自编码器 [论文]	35.42 ± 0.90
均值教师 ResNet-152	9.11 ± 0.12
全部标签，当前最优结果 [论文]	3.79

CIFAR-10 使用 4000 个标签	测试错误率
CT-GAN [论文]	9.98 ± 0.21
均值教师 ResNet-26	6.28 ± 0.15
全部标签，当前最优结果 [论文]	2.86

实现

我们提供了两个实现版本，分别针对 TensorFlow 和 PyTorch。 PyTorch 版本可能更容易根据你的需求进行调整， 因为它遵循典型的 PyTorch 编程习惯，并且可以很方便地集成你的模型和数据集。如果有什么不清楚的地方，请随时联系我。

关于论文中的实验结果，使用传统卷积网络架构的实验是用 TensorFlow 版本完成的。 而使用残差网络的实验则是用 PyToch 版本完成的。

超参数选择及其他调优技巧

均值教师引入了两个新的超参数：EMA 衰减率和一致性损失权重。这两个参数的最佳取值取决于数据集、模型以及迷你批次的组成。此外，你还需要决定如何在迷你批次中混合无标签样本和有标签样本。

以下是一些入门级的经验法则：

• 如果你在处理一个新的数据集，最简单的方法是从仅有有标签数据开始，进行纯监督训练。等你对模型架构和超参数感到满意后，再加入均值教师。通常情况下，同样的网络仍然适用，不过你可能需要适当降低之前在小数据集上使用的正则化强度，比如权重衰减。
• 均值教师需要一定的噪声才能达到最佳效果。实践中，最好的噪声来源往往是随机的输入增强。你可以尝试任何相关的增强技术——算法会训练模型对这些变换保持不变。
• 建议在每个迷你批次中预留一部分用于有标签样本。这样，早期的监督信号会足够强，有助于快速训练并避免陷入不确定性状态。在 PyTorch 的示例中，我们通常将迷你批次的四分之一或二分之一用于有标签样本，其余部分用于无标签样本。（参见 PyTorch 代码中的 TwoStreamBatchSampler。）
• 对于 EMA 衰减率，0.999 通常是一个不错的起点。
• 一致性损失函数可以选择 MSE 或 KL 散度。对于 KL 散度，一致性损失权重一般在 1.0 到 10.0 之间；而对于 MSE，则通常介于类别数和类别数的平方之间。在小规模数据集上，我们发现 MSE 往往能取得更好的效果，不过 KL 散度的表现也一直很稳定。
• 在训练初期的前几个 epoch 中逐步增加一致性损失，直到教师模型开始给出较为准确的预测，可能会有所帮助。
• 我们在 PyTorch 示例中还使用了一个额外的小技巧：在顶层设置两个独立的 logits 层。其中一个用于有标签样本的分类，另一个用于预测教师的输出。然后在这两个预测的 logits 之间再添加一项额外的损失。这一做法的目的与逐步增加一致性损失类似：在训练初期，教师的预测可能还不太准确，因此需要暂时放松分类预测与一致性损失之间的关联。（参见 PyTorch 实现中的 --logit-distance-cost 参数。）

Mean Teacher 快速上手指南

Mean Teacher 是一种高效的半监督学习方法，通过维护一个“教师”模型（学生模型权重的指数移动平均）来指导学生模型的训练，从而在标签数据稀缺的情况下显著提升性能。本项目提供 TensorFlow 和 PyTorch 两种实现，本指南以更易扩展的 PyTorch 版本为例。

环境准备

• 操作系统: Linux, macOS 或 Windows
• Python: 3.6 及以上版本
• 核心依赖:
  • PyTorch (推荐最新稳定版)
  • torchvision
  • NumPy
  • Pillow
• 硬件建议: 推荐使用 NVIDIA GPU 以加速训练过程。

安装步骤

1. 克隆项目仓库

   git clone https://github.com/CuriousAI/mean-teacher.git
   cd mean-teacher/pytorch
   

2. 安装 Python 依赖 建议使用国内镜像源（如清华源）加速安装：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   注：如果 requirements.txt 未明确指定 PyTorch 版本，请根据官方文档单独安装适合你 CUDA 版本的 PyTorch：

   # 示例：安装 CPU 版本 (如有需要)
   pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 示例：安装 CUDA 11.8 版本 (请根据实际情况调整)
   # pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
   

基本使用

以下示例展示如何在 CIFAR-10 数据集上启动半监督训练。该脚本会自动处理数据下载、增强以及 Mean Teacher 的核心逻辑（一致性成本计算与 EMA 更新）。

1. 运行训练脚本 在项目根目录 (pytorch/) 下执行：

   python main.py --dataset cifar10 --labels 4000 --arch resnet26
   

   参数说明：

   • --dataset: 数据集名称 (支持 cifar10, cifar100, mnist 等)。
   • --labels: 使用的有标签样本数量 (例如 4000 表示仅用 4000 张带标签图片，其余作为无标签数据)。
   • --arch: 模型架构 (例如 resnet26, convnet)。

2. 关键超参数调优建议 为了获得最佳效果，可根据实际情况调整以下参数（直接追加到上述命令后）：

   • EMA 衰减率: 默认通常为 0.999，可作为起点。

     --ema-decay 0.999
     

   • 一致性成本权重: 若使用 KL 散度，建议在 1.0 到 10.0 之间；若使用 MSE，通常在类别数到类别数平方之间。

     --consistency 10.0
     

   • Logit 距离成本: 用于在训练初期解耦分类预测与一致性成本，防止教师模型早期预测不准影响训练。

     --logit-distance-cost 0.1
     

   完整示例命令：

   python main.py --dataset cifar10 --labels 4000 --arch resnet26 --ema-decay 0.999 --consistency 10.0 --logit-distance-cost 0.1
   

训练过程中，模型将自动利用无标签数据的一致性约束来优化学生网络，并同步更新教师网络的权重。