CLIP损失函数在跨模态检索中的实战优化：从理论到PyTorch实现

1次阅读

没有评论

共计 2225 个字符，预计需要花费 6 分钟才能阅读完成。

传统跨模态检索面临图文特征空间不一致的难题：

图像和文本的原始特征分布差异大，直接计算相似度效果差
人工设计的损失函数（如三元组损失）难以捕捉模态间复杂关系
当 batch size 超过 512 时，CLIP 原始实现会出现显存不足问题（24GB 显存仅支持约 640 批次）

对称 InfoNCE 损失定义为：

$$
\mathcal{L}{i→j} = -\log\frac{\exp(s
$$}/\tau)}{\sum_{k=1}^N\exp(s_{i,k}/\tau)

$$
\mathcal{L}{j→i} = -\log\frac{\exp(s
$$}/\tau)}{\sum_{k=1}^N\exp(s_{k,i}/\tau)

最终损失为两个方向的均值：
$$
\mathcal{L} = \frac{1}{2N}\sum_{i=1}^N(\mathcal{L}{i→j} + \mathcal{L})
$$

梯度累积 ：将大 batch 拆分为多个 micro-batch，显存需求降低为原来的 1 / 累积步长
混合精度训练 ：
前向计算使用 FP16
损失计算保持 FP32 防数值溢出
梯度更新转回 FP32
分布式负样本 ：通过 all_gather 收集全局负样本，使每个 GPU 获得完整负样本库

import torch
import torch.distributed as dist
from torch.cuda.amp import autocast

class CLIPLoss(torch.nn.Module):
    def __init__(self, temp_init=0.07):
        super().__init__()
        # 可学习温度参数（log 域防止负值）self.logit_scale = torch.nn.Parameter(torch.log(torch.tensor(1/temp_init)))

    def forward(self, image_feat, text_feat):
        device = image_feat.device

        # 归一化特征
        image_feat = torch.nn.functional.normalize(image_feat, dim=-1)
        text_feat = torch.nn.functional.normalize(text_feat, dim=-1)

        # 分布式特征收集
        if dist.is_initialized():
            image_feat_all = [torch.zeros_like(image_feat) for _ in range(dist.get_world_size())]
            text_feat_all = [torch.zeros_like(text_feat) for _ in range(dist.get_world_size())]
            dist.all_gather(image_feat_all, image_feat)
            dist.all_gather(text_feat_all, text_feat)
            image_feat_all = torch.cat(image_feat_all)
            text_feat_all = torch.cat(text_feat_all)
        else:
            image_feat_all = image_feat
            text_feat_all = text_feat

        # 混合精度计算
        with autocast():
            logit_scale = torch.clamp(self.logit_scale.exp(), max=100)
            logits = logit_scale * image_feat @ text_feat_all.T

            labels = torch.arange(len(logits), device=device)
            loss_i = torch.nn.functional.cross_entropy(logits, labels)
            loss_t = torch.nn.functional.cross_entropy(logits.T, labels)

        return (loss_i + loss_t) / 2

关键实现细节：