TENTTTA希望利用测试阶段获得的目标域数据让模型对其进行适应用适应后的模型预测。但测试阶段没有真实标签那应该用什么损失函数更新模型1.Tent优化目标是最小化预测熵当预测集中在某个类别时熵较低。Tent优化目标是让模型对目标域数据的预测更加确定。2.只更新BN的仿射参数核心过程# 1. 使用目标域数据预测 logits model(target_x) # 2. 计算预测熵 probs torch.softmax(logits, dim1) entropy -( probs * torch.log(probs.clamp_min(1e-8)) ).sum(dim1).mean() # 3. 计算梯度 entropy.backward() # 4. 只更新 BN affine 参数 optimizer.step()关于BN层最常见的CNN结构输入图像│▼Conv2d卷积层 提取图像特征│▼BatchNorm2dBN 层对中间特征进行归一化和仿射变换│▼ReLU激活函数引入非线性│▼Conv2d│▼BatchNorm2d│▼ReLU│▼...│▼全连接层│▼分类结果神经网络每一层接收到的数据其数值范围可能非常不同而且会随着训练不断变化。如果中间特征的数值尺度变化很大神经网络的训练可能变得困难。引入归一化就可以把输入调整到相近的数值尺度减轻由于输入尺度差异造成的优化困难。BN层作用是对中间特征进行归一化和仿射变换不会改变特征图的形状而是改变特征值的分布。简单代码 TENT: Fully Test-Time Adaptation by Entropy Minimization 数据集MNIST 模型简单 CNN BatchNorm 域偏移Gaussian Noise 比较 1. Source Model 2. BN Adapt 3. TENT 运行 pip install torch torchvision matplotlib python tent_mnist.py import copy import random import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms # # 1. 配置 # SEED 42 BATCH_SIZE 128 TRAIN_EPOCHS 3 TRAIN_LR 0.001 TENT_LR 0.001 TENT_STEPS 1 NOISE_STD 0.8 DEVICE torch.device( cuda if torch.cuda.is_available() else mps if torch.backends.mps.is_available() else cpu ) def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) set_seed(SEED) print(Device:, DEVICE) # # 2. 数据集 # # 源域正常 MNIST source_transform transforms.ToTensor() # 目标域添加高斯噪声 class AddGaussianNoise: def __init__(self, std0.8): self.std std def __call__(self, x): noise torch.randn_like(x) * self.std x x noise return torch.clamp(x, 0.0, 1.0) # 给正常的 MNIST 图片添加高斯噪声人为模拟 Distribution Shift分布偏移。 target_transform transforms.Compose([# 把多个数据处理操作按照指定顺序组合起来依次执行。 transforms.ToTensor(), AddGaussianNoise(stdNOISE_STD) ]) train_dataset datasets.MNIST( root./data, trainTrue, downloadTrue, transformsource_transform ) clean_test_dataset datasets.MNIST( root./data, trainFalse, downloadTrue, transformsource_transform ) target_test_dataset datasets.MNIST( root./data, trainFalse, downloadTrue, transformtarget_transform ) train_loader DataLoader( train_dataset, batch_sizeBATCH_SIZE, shuffleTrue ) clean_test_loader DataLoader( clean_test_dataset, batch_sizeBATCH_SIZE, shuffleFalse ) # 为了公平比较预先固定目标域的噪声样本。 # 否则三种方法每次读取时会生成不同噪声。 fixed_target_images [] fixed_target_labels [] for x, y in DataLoader( target_test_dataset, batch_sizeBATCH_SIZE, shuffleFalse ): fixed_target_images.append(x) fixed_target_labels.append(y) fixed_target_images torch.cat(fixed_target_images) fixed_target_labels torch.cat(fixed_target_labels) target_loader DataLoader( torch.utils.data.TensorDataset( fixed_target_images, fixed_target_labels ), batch_sizeBATCH_SIZE, shuffleFalse ) # # 3. CNN 模型 # class CNN(nn.Module): def __init__(self): super().__init__() # 提取图片特征 self.features nn.Sequential( nn.Conv2d(1, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2) ) # 根据特征进行分类 self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * 7 * 7, 128), nn.ReLU(), nn.Linear(128, 10) ) # 定义数据如何经过网络 def forward(self, x): x self.features(x) x self.classifier(x) return x # # 4. 源域监督训练 # # 训练源域模型把初始的 CNN 模型放到干净的 MNIST 训练集上进行有监督训练 def train_source(model, loader): # 将模型切换到训练模式 model.train() # 创建Adam优化器 optimizer torch.optim.Adam( model.parameters(), # 源域训练更新整个 CNNTENT 只更新 BN 的仿射参数 lrTRAIN_LR ) for epoch in range(TRAIN_EPOCHS): # 一个 Epoch 完整遍历一次训练集 total_loss 0.0 correct 0 total 0 # 每执行一次循环DataLoader 自动返回一个 Batch for x, y in loader: x x.to(DEVICE) # 当前 Batch 的 128 张图片 y y.to(DEVICE) # 这 128 张图片对应的标签 optimizer.zero_grad() # 前向传播logits存模型输出的原始分类分数 logits model(x) # 计算交叉熵损失 loss F.cross_entropy(logits, y) loss.backward() # 更新参数 optimizer.step() # 累加每个 Batch 的 Loss total_loss loss.item() # 将 Logits 转换为预测类别 predictions logits.argmax(dim1) # 统计正确数量 correct (predictions y).sum().item() total y.size(0) accuracy correct / total * 100 print( fEpoch {epoch 1}/{TRAIN_EPOCHS} | fLoss: {total_loss / len(loader):.4f} | fAccuracy: {accuracy:.2f}% ) # # 5. 普通模型评估 # # 使用已经训练好的 CNN 对数据集进行预测并计算分类准确率 torch.no_grad() # 执行整个 evaluate() 函数时不记录用于反向传播的计算图 def evaluate(model, loader): model.eval() # 遍历数据集并预测 correct 0 total 0 for x, y in loader: x x.to(DEVICE) y y.to(DEVICE) logits model(x) predictions logits.argmax(dim1) correct (predictions y).sum().item() total y.size(0) return correct / total * 100 # # 6. 配置 BN Adapt # def configure_bn_adapt(model): # 其他层保持 eval避免无关模块改变行为 model.eval() for module in model.modules(): if isinstance(module, nn.BatchNorm2d): # 只有 BN 层才会执行下面的操作其他层不受影响。 # BN 使用当前测试 Batch 的统计量 module.train() # 不维护 running statistics # 不再跟踪和累计这些统计量 module.track_running_stats False # 把原来保存的均值和方差清空 module.running_mean None module.running_var None return model # # 7. BN Adapt 评估 # torch.no_grad() def evaluate_bn_adapt(model, loader): configure_bn_adapt(model) correct 0 total 0 for x, y in loader: x x.to(DEVICE) y y.to(DEVICE) logits model(x) predictions logits.argmax(dim1) correct (predictions y).sum().item() total y.size(0) return correct / total * 100 # # 8. 配置 TENT # def configure_tent(model): # 先将整个模型设置为 eval model.eval() # 冻结所有参数 for param in model.parameters(): param.requires_grad_(False) tent_params [] for module in model.modules(): if isinstance(module, nn.BatchNorm2d): # BN 使用测试批次统计量 module.train() module.track_running_stats False module.running_mean None module.running_var None # 只允许 BN affine 参数更新 module.weight.requires_grad_(True) module.bias.requires_grad_(True) tent_params.append(module.weight) tent_params.append(module.bias) return model, tent_params # # 9. 预测熵 # def softmax_entropy(logits): log_probs F.log_softmax(logits, dim1) probs F.softmax(logits, dim1) entropy -(probs * log_probs).sum(dim1) return entropy # # 10. TENT 单步更新 # def tent_step(model, x, optimizer): optimizer.zero_grad() # Step 1前向传播 logits model(x) # Step 2计算每个样本的预测熵 entropy softmax_entropy(logits) # Step 3计算 Batch 平均熵 loss entropy.mean() # Step 4反向传播 loss.backward() # Step 5更新 BN affine optimizer.step() return loss.item() # # 11. TENT 测试 # def evaluate_tent(model, loader): model, tent_params configure_tent(model) optimizer torch.optim.SGD( tent_params, lrTENT_LR ) correct 0 total 0 entropy_history [] for x, y in loader: x x.to(DEVICE) y y.to(DEVICE) # 在当前 Batch 上进行无监督适应 for _ in range(TENT_STEPS): entropy tent_step( model, x, optimizer ) entropy_history.append(entropy) # 使用更新后的参数重新预测当前 Batch with torch.no_grad(): logits model(x) predictions logits.argmax(dim1) correct (predictions y).sum().item() total y.size(0) accuracy correct / total * 100 return accuracy, entropy_history # # 12. 主程序 # def main(): print(\n Source Training ) model CNN().to(DEVICE) train_source(model, train_loader) # 保存源模型 torch.save( model.state_dict(), source_mnist.pth ) # -------------------------------------------------------- # Source Model # -------------------------------------------------------- print(\n Source Evaluation ) clean_acc evaluate( model, clean_test_loader ) source_target_acc evaluate( model, target_loader ) print(fClean Accuracy: {clean_acc:.2f}%) print(fTarget Accuracy: {source_target_acc:.2f}%) # -------------------------------------------------------- # BN Adapt # -------------------------------------------------------- print(\n BN Adapt ) bn_model copy.deepcopy(model) bn_acc evaluate_bn_adapt( bn_model, target_loader ) print(fBN Adapt Accuracy: {bn_acc:.2f}%) # -------------------------------------------------------- # TENT # -------------------------------------------------------- print(\n TENT ) tent_model copy.deepcopy(model) tent_acc, entropy_history evaluate_tent( tent_model, target_loader ) print(fTENT Accuracy: {tent_acc:.2f}%) # -------------------------------------------------------- # 结果汇总 # -------------------------------------------------------- print(\n Results ) print(fClean Source: {clean_acc:.2f}%) print(fCorrupted Source: {source_target_acc:.2f}%) print(fBN Adapt: {bn_acc:.2f}%) print(fTENT: {tent_acc:.2f}%) print( fTENT vs Source: f{tent_acc - source_target_acc:.2f} pp ) print( fTENT vs BN Adapt: f{tent_acc - bn_acc:.2f} pp ) if __name__ __main__: main()