TTA详细学习

TENT

TTA希望利用测试阶段获得的目标域数据让模型对其进行适应，用适应后的模型预测。但测试阶段没有真实标签，那应该用什么损失函数更新模型？

1.Tent优化目标是最小化预测熵

当预测集中在某个类别时，熵较低。Tent优化目标是让模型对目标域数据的预测更加确定。

2.只更新BN的仿射参数

核心过程

# 1. 使用目标域数据预测 logits = model(target_x) # 2. 计算预测熵 probs = torch.softmax(logits, dim=1) entropy = -( probs * torch.log(probs.clamp_min(1e-8)) ).sum(dim=1).mean() # 3. 计算梯度 entropy.backward() # 4. 只更新 BN affine 参数 optimizer.step()

关于BN层

最常见的CNN结构

输入图像
│
▼
Conv2d（卷积层） 提取图像特征
│
▼
BatchNorm2d（BN 层）对中间特征进行归一化和仿射变换
│
▼
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, std=0.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(std=NOISE_STD) ]) train_dataset = datasets.MNIST( root="./data", train=True, download=True, transform=source_transform ) clean_test_dataset = datasets.MNIST( root="./data", train=False, download=True, transform=source_transform ) target_test_dataset = datasets.MNIST( root="./data", train=False, download=True, transform=target_transform ) train_loader = DataLoader( train_dataset, batch_size=BATCH_SIZE, shuffle=True ) clean_test_loader = DataLoader( clean_test_dataset, batch_size=BATCH_SIZE, shuffle=False ) # 为了公平比较，预先固定目标域的噪声样本。 # 否则三种方法每次读取时会生成不同噪声。 fixed_target_images = [] fixed_target_labels = [] for x, y in DataLoader( target_test_dataset, batch_size=BATCH_SIZE, shuffle=False ): 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_size=BATCH_SIZE, shuffle=False ) # ============================================================ # 3. CNN 模型 # ============================================================ class CNN(nn.Module): def __init__(self): super().__init__() # 提取图片特征 self.features = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size=3, padding=1), 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(), # 源域训练更新整个 CNN，TENT 只更新 BN 的仿射参数 lr=TRAIN_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(dim=1) # 统计正确数量 correct += (predictions == y).sum().item() total += y.size(0) accuracy = correct / total * 100 print( f"Epoch {epoch + 1}/{TRAIN_EPOCHS} | " f"Loss: {total_loss / len(loader):.4f} | " f"Accuracy: {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(dim=1) 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(dim=1) 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, dim=1) probs = F.softmax(logits, dim=1) entropy = -(probs * log_probs).sum(dim=1) 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, lr=TENT_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(dim=1) 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(f"Clean Accuracy: {clean_acc:.2f}%") print(f"Target 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(f"BN 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(f"TENT Accuracy: {tent_acc:.2f}%") # -------------------------------------------------------- # 结果汇总 # -------------------------------------------------------- print("\n========== Results ==========") print(f"Clean Source: {clean_acc:.2f}%") print(f"Corrupted Source: {source_target_acc:.2f}%") print(f"BN Adapt: {bn_acc:.2f}%") print(f"TENT: {tent_acc:.2f}%") print( f"TENT vs Source: " f"{tent_acc - source_target_acc:+.2f} pp" ) print( f"TENT vs BN Adapt: " f"{tent_acc - bn_acc:+.2f} pp" ) if __name__ == "__main__": main()