如何在多标签分类任务中使用Torchvision MixUp？

原torchvision.transforms.v2.MixUp仅针对单标签分类任务设计（标签为形状(batch_size,)的类别索引），无法直接处理[0.0, 1.0, 0.0, 1.0]这类形状为(batch_size, num_classes)的多标签向量。你之前用argmax转单索引的方式会丢失多标签信息，以下是两种可行的修改方案：

方案一：继承原MixUp类重写标签逻辑 #

直接复用原MixUp的图像混合逻辑，仅修改标签的处理方式——对多标签向量做线性混合，而非转成one-hot编码：

import torch
from torchvision.transforms.v2 import MixUp

class MultiLabelMixUp(MixUp):
    def _transform(self, inpt, label, params):
        lam = params["lam"]
        # 混合图像（和原MixUp逻辑一致）
        if isinstance(inpt, dict):
            for key in inpt:
                inpt[key] = lam * inpt[key] + (1 - lam) * inpt[key].flip(0)
        else:
            inpt = lam * inpt + (1 - lam) * inpt.flip(0)
        
        # 混合多标签：直接对标签向量做线性加权
        if label is not None:
            label = lam * label + (1 - lam) * label.flip(0)
        
        return inpt, label

使用方式 #

和原MixUp完全一致，只需传入形状为(batch_size, num_classes)的多标签张量：

mixup_transform = MultiLabelMixUp(alpha=0.8)
# 假设images形状为(batch, 3, H, W)，labels形状为(batch, num_classes)
mixed_images, mixed_labels = mixup_transform(images, labels)

方案二：手动实现多标签MixUp逻辑 #

如果不需要依赖原MixUp类，可以手动实现更轻量化的版本，逻辑更直观：

import torch

def multi_label_mixup(images, labels, alpha=0.8):
    # 生成Beta分布的混合权重
    lam = torch.distributions.Beta(alpha, alpha).sample((images.size(0),)).to(images.device)
    # 扩展维度以匹配图像和标签的形状（广播计算）
    img_lam = lam.view(-1, 1, 1, 1)  # 适配图像形状(batch, channels, H, W)
    label_lam = lam.view(-1, 1)       # 适配标签形状(batch, num_classes)
    
    # 生成打乱的batch索引，用于取配对样本
    shuffle_idx = torch.randperm(images.size(0)).to(images.device)
    
    # 执行图像和标签的混合
    mixed_images = img_lam * images + (1 - img_lam) * images[shuffle_idx]
    mixed_labels = label_lam * labels + (1 - label_lam) * labels[shuffle_idx]
    
    return mixed_images, mixed_labels

使用方式 #

在训练循环或数据增强环节直接调用：

# 假设images、labels为当前batch的张量
mixed_images, mixed_labels = multi_label_mixup(images, labels, alpha=0.8)

关键逻辑说明 #

多标签任务的MixUp核心是保留所有类别的标签信息：直接对多标签向量做线性加权混合，而非将标签压缩为单类别索引。这样混合后的标签依然是(batch_size, num_classes)的向量，每个维度对应类别的置信度，完全适配多标签分类的损失计算（比如BCELoss）。

内容的提问来源于stack exchange，提问作者will