AI自监督学习：无标签不够用了之后

自监督学习 对比学习 预训练 PyTorch 表示学习
作者: 醉月思 发布日期: 2022-04-06 分类: 技术实践

前阵子AI自监督学习连续改了几轮，趁还记得写成备忘。

场景和约束

项目背景：工业质检场景的缺陷检测，要识别十几种不同类型的缺陷。

约束条件很现实：

  • 数据总量约 30 万张图像，但标注预算只够标 2 万张
  • 计算资源有限，两块 RTX 3090，不是 A100 那种豪配置
  • 时间紧迫，三个月要出个能用的模型
  • 后期需要持续添加新缺陷类型，模型要能快速适配

一开始想的方案很直接：直接用 2 万张标注数据训练个 ResNet-50。效果勉强能用，准确率 78% 左右，但泛化能力很差，换一批生产环境的数据，准确率直接掉到 60%。后来分析原因：2 万张数据对于十几个类别来说，每个类别平均也就一两千张，训练集和测试集的分布差异又比较大，模型过度拟合了训练集的特定分布。

自监督学习的选择

自监督学习的方法挺多，但真正适合这种资源受限的场景，其实就几个路子：

  1. 对比学习：让模型学会"相似和不相似"这个概念，代表是 SimCLR、MoCo
  2. 掩码图像建模：让模型学会"补全被遮挡的部分"，代表是 MAE
  3. 自回归建模：让模型学会"预测下一个 patch"，代表是 BEiT

考虑到计算资源和时间限制，最终选择了 SimCLR。理由很简单：实现相对简单，不需要复杂的架构改动，而且 PyTorch 生态里有现成的库，上手成本低。

SimCLR 的核心思路：对同一张图像做两次不同的数据增强，得到两个"视图"，让模型学会把这两个视图识别为"相似"，把不同图像的视图识别为"不相似"。

环境和工具链

这次实践用的工具链很常见，但版本组合确实踩过坑：

# 环境配置
python 3.10
pytorch 2.1.0
torchvision 0.16.0
timm 0.9.12  # 用于预训练模型
pytorch-lightning 2.1.0  # 训练框架
tensorboard 2.15.0  # 可视化

一开始用 PyTorch 2.0，发现在某些混合精度训练场景下会有 NaN 梯度的问题，升级到 2.1.0 后好了。timm 这个库很实用，各种预训练模型直接调，不用自己写 ResNet 的实现。

SimCLR 的实现

先写个最简单的数据增强 pipeline，这是 SimCLR 的关键：

# data/augmentation.py
import torchvision.transforms as T
from torchvision.transforms import InterpolationMode

class SimCLRTransform:
    """SimCLR 的数据增强组合"""

    def __init__(self, size=224, s=1.0):
        color_jitter = T.ColorJitter(
            0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s
        )
        self.transform = T.Compose([
            T.RandomResizedCrop(size, scale=(0.2, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.RandomApply([color_jitter], p=0.8),
            T.RandomGrayscale(p=0.2),
            T.GaussianBlur(kernel_size=int(0.1 * size), sigma=(0.1, 2.0)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, x):
        # 对同一张图像生成两个不同的增强视图
        return self.transform(x), self.transform(x)

这里有个坑：RandomResizedCrop 的 scale=(0.2, 1.0) 意味着图像可能会被裁剪到原来的 20%，这个比例对于缺陷检测场景来说有点激进。后来调整到 (0.5, 1.0)，保留更多图像原始信息，效果反而更好。

然后是模型和投影头：

# models/simclr.py
import torch
import torch.nn as nn
import timm

class SimCLRModel(nn.Module):
    """SimCLR 模型：编码器 + 投影头"""

    def __init__(self, backbone='resnet50', projection_dim=128):
        super().__init__()
        # 使用 timm 的预训练模型作为编码器
        self.encoder = timm.create_model(
            backbone,
            pretrained=False,  # 我们要用自己的数据预训练
            num_classes=0  # 移除分类头
        )

        # 投影头：2 层 MLP
        self.projection_head = nn.Sequential(
            nn.Linear(self.encoder.num_features, 512),
            nn.ReLU(),
            nn.Linear(512, projection_dim)
        )

    def forward(self, x):
        features = self.encoder(x)
        projections = self.projection_head(features)
        return projections

投影头是个关键点。一开始觉得"既然要学表示，为什么不直接用编码器的特征"，后来试过，效果确实差不少。原因大概是：投影头把特征映射到一个"更适合对比学习"的空间，去掉了一些对对比任务没用但对下游任务有用的信息。

损失函数用的是 NT-Xent（Normalized Temperature-scaled Cross Entropy Loss）：

# losses/nt_xent.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class NTXentLoss(nn.Module):
    """SimCLR 的损失函数"""

    def __init__(self, temperature=0.5):
        super().__init__()
        self.temperature = temperature
        self.criterion = nn.CrossEntropyLoss(reduction="sum")

    def forward(self, z_i, z_j):
        """
        z_i: (batch_size, projection_dim)
        z_j: (batch_size, projection_dim)
        """
        batch_size = z_i.shape[0]

        # 拼接正样本对
        z = torch.cat((z_i, z_j), dim=0)

        # 计算相似度矩阵
        sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2)

        # 除以温度参数
        sim = sim / self.temperature

        # 标签：正样本对在相似度矩阵中的位置
        labels = torch.arange(batch_size, device=z.device)
        labels = torch.cat((labels + batch_size, labels), dim=0)

        # 计算损失
        loss = self.criterion(sim, labels)

        return loss / (2 * batch_size)

这里有个细节：temperature 参数的选择。一开始用 0.1，模型很快就"记住"了训练数据，泛化能力很差；后来调到 0.5，情况好不少。这个参数确实需要根据具体场景调，没有固定值。

训练过程和踩过的坑

训练过程本身不复杂，但踩的坑不少。

1. Batch Size 的选择

SimCLR 对 batch size 很敏感，论文里建议用 4096，但两块 3090 根本跑不动。试了几个方案：

  • 256：训练不稳定，损失震荡厉害
  • 512：稍微好点，但仍然不够稳
  • 1024：勉强稳定，但显存接近极限，容易 OOM

最后选了 1024，配合梯度累积：

# trainer.py
import pytorch_lightning as pl

class SimCLRTrainer(pl.LightningModule):
    def __init__(self, model, learning_rate=1e-3, batch_size=1024):
        super().__init__()
        self.model = model
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.criterion = NTXentLoss(temperature=0.5)

    def training_step(self, batch, batch_idx):
        (x_i, x_j), _ = batch  # 我们不需要标签
        z_i = self.model(x_i)
        z_j = self.model(x_j)
        loss = self.criterion(z_i, z_j)

        self.log('train_loss', loss, prog_bar=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=self.learning_rate,
            weight_decay=1e-4
        )
        return optimizer

# 实际训练
trainer = pl.Trainer(
    max_epochs=100,
    accumulate_grad_batches=4,  # 模拟 batch size 4096
    precision='16-mixed',  # 混合精度训练
    gradient_clip_val=1.0  # 梯度裁剪
)

梯度累积确实有效，但训练速度慢了很多。100 个 epoch 跑了大概三天，这个时间成本确实是需要考虑的。

2. 学习率调度

一开始用固定的学习率，模型训练到 60 个 epoch 左右就停滞了。后来加了个 cosine 学习率调度：

# 学习率调度器
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=100,  # 总 epoch 数
    eta_min=1e-5  # 最小学习率
)

效果确实有提升，但提升不算巨大。后来想想，可能数据量本身就不算很大，学习率的影响没论文里说得那么明显。

3. 混合精度训练的坑

用 precision='16-mixed' 的时候，遇到一个问题：某些 batch 的 loss 突然变成 NaN。排查了一下午，发现是投影头的最后一层权重在某些极端情况下会变成 0，导致除零。

解决方法是加一个小的 epsilon：

# 修改后的投影头
self.projection_head = nn.Sequential(
    nn.Linear(self.encoder.num_features, 512),
    nn.ReLU(),
    nn.Linear(512, projection_dim),
    nn.BatchNorm1d(projection_dim)  # 加一个 BN 层稳定训练
)

加了 BN 层后，NaN 问题基本消失了。这个套路后来在别的项目里也用得上。

4. 早停策略的权衡

自监督学习没有明确的"验证集准确率"这种指标，怎么判断"可以停了"是个问题。一开始用训练 loss，但 loss 稳定下降不代表学到的东西有用。

后来想了个办法：每个 epoch 保存一次模型，然后用 2 万张标注数据快速评估一下"特征质量"：

def evaluate_feature_quality(model, labeled_dataloader):
    """评估特征质量：用一个小分类器测试特征的区分度"""
    model.eval()
    features = []
    labels = []

    with torch.no_grad():
        for images, target in labeled_dataloader:
            feat = model.encoder(images)  # 只用编码器
            features.append(feat)
            labels.append(target)

    features = torch.cat(features, dim=0).cpu().numpy()
    labels = torch.cat(labels, dim=0).cpu().numpy()

    # 用简单的线性分类器测试
    from sklearn.linear_model import LogisticRegression
    clf = LogisticRegression(max_iter=1000)
    clf.fit(features[:10000], labels[:10000])  # 用一半数据训练
    accuracy = clf.score(features[10000:], labels[10000:])  # 用另一半数据测试

    return accuracy

这个评估方法不算精确，但能给出一个"大概"的趋势。如果特征质量连续 5 个 epoch 没提升，就可以考虑停了。

从预训练到微调

自监督训练完成后，下一步是微调。这个过程相对简单，但也有一些技巧。

1. 微调策略的选择

有几种策略：

  • 冻结编码器，只训练分类头：速度快，但效果可能受限
  • 解冻最后几层， jointly 训练：平衡效果和速度
  • 全部解冻，用小学习率微调：效果最好，但耗时最长

试了一圈，发现策略 2 最适合我们的场景：

# 微调模型
class FineTunedModel(nn.Module):
    def __init__(self, pretrained_encoder, num_classes=13):
        super().__init__()
        self.encoder = pretrained_encoder

        # 冻结编码器的前几层
        for param in self.encoder.parameters():
            param.requires_grad = False

        # 解冻最后两层
        for name, param in self.encoder.named_parameters():
            if 'layer4' in name or 'layer3' in name:
                param.requires_grad = True

        # 新的分类头
        self.classifier = nn.Sequential(
            nn.Linear(self.encoder.num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        features = self.encoder(x)
        return self.classifier(features)

2. 学习率的差异

微调的学习率要比预训练小很多，而且编码器和分类头的学习率要分开设置：

# 分类头用大学习率，编码器用小学习率
optimizer = torch.optim.AdamW([
    {'params': model.encoder.parameters(), 'lr': 1e-4},  # 编码器
    {'params': model.classifier.parameters(), 'lr': 1e-3}  # 分类头
])

这个套路在迁移学习里很常见，但确实有用。一开始用统一学习率，编码器被"破坏"得厉害，预训练的好处基本没了。

3. 数据增强的调整

微调阶段的数据增强和预训练阶段不太一样。预训练阶段用的增强很激进，目的是让模型学到"不变性"；微调阶段需要保守一点，保留更多"任务相关"的信息。

# 微调阶段的数据增强
finetune_transform = T.Compose([
    T.RandomResizedCrop(224, scale=(0.8, 1.0)),  # 缩小裁剪范围
    T.RandomHorizontalFlip(p=0.5),
    T.RandomApply([T.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.5),  # 减弱颜色抖动
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

调整后，微调的收敛速度快了不少，最终准确率也比用预训练增强方案高 3 个百分点左右。

效果对比

最后对比了几种方案的效果（准确率是在独立测试集上的）：

方案训练数据测试集准确率训练时间
直接训练2 万标注数据78.2%8 小时
ImageNet 预训练 + 微调2 万标注数据82.1%2 小时
自监督预训练 + 微调2 万标注数据 + 28 万无标签数据85.6%3 天预训练 + 6 小时微调

三种方案最终测试准确率差距一目了然，尤其能看出自监督预训练相对 ImageNet 迁移的额外收益。

直接训练、ImageNet 预训练与 SimCLR 自监督三种方案的测试集准确率对比

自监督预训练在领域差异较大的工业质检场景里，比 ImageNet 迁移再多出约 3.5 个百分点，且对新缺陷类别适配更好。

几个观察：

  1. ImageNet 预训练确实有用，但对于工业质检这种领域差异很大的场景，收益有限
  2. 自监督预训练虽然耗时，但效果明显比直接训练好，也比 ImageNet 预训练好
  3. 训练时间确实是个问题，但如果"只预训练一次，长期复用"，这个成本可以摊薄

还有一个意外的发现：自监督预训练对"新缺陷类别"的适应性更好。后来加了 3 个新的缺陷类别，用原来的模型微调，准确率很快就能到 80% 以上。而直接训练的模型，新类别的准确率一直在 60% 左右徘徊。

其他尝试和坑

MoCo 的尝试

在 SimCLR 之外，还试过 MoCo（Momentum Contrast）。MoCo 的优势是不需要大 batch size，理论上更适合资源受限的场景。

# MoCo 的核心：队列和动量编码器
class MoCoModel(nn.Module):
    def __init__(self, backbone='resnet50', dim=128, K=65536, m=0.999, T=0.07):
        super().__init__()

        # 编码器 q
        self.encoder_q = timm.create_model(backbone, pretrained=False, num_classes=dim)

        # 编码器 k（动量更新）
        self.encoder_k = timm.create_model(backbone, pretrained=False, num_classes=dim)

        # 初始化 k 的权重为 q
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data.copy_(param_q.data)
            param_k.requires_grad = False  # k 不需要梯度

        self.queue = torch.randn(dim, K)
        self.queue = nn.functional.normalize(self.queue, dim=0)
        self.register_buffer("queue", self.queue)

        self.K = K
        self.m = m
        self.T = T

    @torch.no_grad()
    def _momentum_update_key_encoder(self):
        """动量更新 k"""
        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys):
        """更新队列"""
        batch_size = keys.shape[0]
        ptr = int(self.queue_ptr)
        assert self.K % batch_size == 0

        # 替换队列中的 keys
        self.queue[:, ptr:ptr + batch_size] = keys.T
        self.queue_ptr = (ptr + batch_size) % self.K

但实际效果不如 SimCLR。可能的原因：

  1. 队列的维护增加了训练的复杂性，调参空间变大
  2. MoCo 对动量参数、队列大小等超参数很敏感，调得不对就不稳
  3. 资源有限的情况下，SimCLR + 梯度累积的方案其实更稳定

掩码图像建模的尝试

还试过简化的 MAE（Masked AutoEncoder），想法是让模型学会"补全被遮挡的部分"，这对缺陷检测场景似乎很合理——缺陷不就是"不正常的补全"吗？

# 简化的 MAE 结构
class SimpleMAE(nn.Module):
    def __init__(self, backbone='resnet50', mask_ratio=0.75):
        super().__init__()
        self.encoder = timm.create_model(backbone, pretrained=False, num_classes=0)
        self.decoder = nn.Sequential(
            nn.Linear(self.encoder.num_features, 512),
            nn.ReLU(),
            nn.Linear(512, 3 * 32 * 32)  # 假设 patch size 32x32
        )
        self.mask_ratio = mask_ratio

    def random_masking(self, x):
        """随机遮挡部分 patch"""
        batch_size, _, _, _ = x.shape
        num_patches = (x.shape[2] // 32) * (x.shape[3] // 32)
        num_mask = int(num_patches * self.mask_ratio)
        # 简化实现：直接随机遮挡
        mask = torch.rand(batch_size, num_patches) > self.mask_ratio
        return mask

    def forward(self, x):
        mask = self.random_masking(x)
        # 这里省略了 patch embedding 和位置编码的细节
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded, mask

但效果不如预期。可能的原因：

  1. 工业质检的图像噪声比较多，简单的"补全"任务学不到有用的特征
  2. 缺陷检测需要的是"异常检测"的能力，而 MAE 学的是"正常模式"
  3. 我们的实现太简化，没有真正用到 patch embedding 和位置编码

这条路后来就没继续走下去，但感觉如果数据更干净一些，MAE 的思路还是有潜力的。

实践经验总结

这次折腾下来，几点比较实在的经验：

  1. 资源受限时，优先考虑简单的方案：SimCLR 比 MoCo、MAE 都简单，调参空间小，不容易踩大坑
  2. 数据增强要针对场景调整：工业质检的图像和自然图像不一样，太激进的增强可能适得其反
  3. 评估指标很重要：没有标注数据时，要想办法评估"特征质量"，不能只看训练 loss
  4. 预训练的成本要摊薄：一次预训练，多次微调，这个经济账才划算
  5. 新类别适配能力是个隐藏优势：自监督预训练的模型对新的缺陷类别适应性更好，这一点后期确实省了不少事

还有一些没完全解决的问题：

  1. 训练时间还是太长，如果有更多 GPU 资源，可以尝试并行训练
  2. 超参数的调优仍然比较依赖经验，自动化调参在自监督学习里不好做
  3. 对数据分布的变化敏感，如果生产环境的图像特征变化很大，可能需要重新预训练

收尾

这次实践最直观的感受是：自监督学习不是"银弹"，但在"数据很多但标注很少"的场景下，确实是个不错的工具。它把"无标签数据"变成了"有用的表示"，然后在标注数据上快速适配，这个思路在很多场景都能用得上。

写这篇文章的时候，模型已经在生产环境跑了半年，准确率稳定在 85% 左右，新缺陷类别的适配也快了不少。虽然预训练花了三天时间，但这个投入长期来看是值得的。

当然，如果你的场景是"数据少、标注多、算力足"，那直接训练或者用 ImageNet 预训练可能更合适。技术方案的选择，最终还是看场景和约束。

代码和配置都是基于实际项目改写，如果发现哪里不对，可能是当时环境和版本不太一样。具体项目用的是 PyTorch 2.1.0、timm 0.9.12、PyTorch Lightning 2.1.0，在两块 RTX 3090 上跑的。如果你也在做类似的项目，希望这篇文章能少踩几个坑。

版权声明: 本文首发于 指尖魔法屋-AI自监督学习：无标签不够用了之后（https://blog.thinkmoon.cn/post/222-self-supervised-learning-representation-from-no-label/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章