TTA详细学习

想七想八不如114082026-10-24 13:09

TENT

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

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

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

2.只更新BN的仿射参数

核心过程

python 复制代码 
# 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层作用是对中间特征进行归一化和仿射变换，不会改变特征图的形状，而是改变特征值的分布。

简单代码

python 复制代码 
"""
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()
相关推荐
探索者轻舞2 小时前
Logprobs 零成本测谎仪——从分母抹杀理论到软预警协同裁决
人工智能·深度学习
数聚天成DeepSData2 小时前
Landsat 影像去哪里下载：选景、批量与脚本检索怎么选
人工智能·深度学习·机器学习·数据集·deepsdata
小蒋观天下2 小时前
【无标题】红外补光近距过曝、远距欠曝：导致车牌AI识别失效调试
人工智能·深度学习·数码相机·算法·机器学习·ai大模型
智码看视界2 小时前
深度学习：迁移学习与预训练模型应用实践
深度学习·微调·迁移学习·小样本·特征提取·#预训练模型·领域适配
段一凡-华北理工大学3 小时前
高炉炼铁机器视觉与智能识别十八讲~系列文章17：实战案例：一个风口监控系统的完整搭建
人工智能·pytorch·深度学习·机器学习·计算机视觉·高炉风口监测系统
哦哦～9213 小时前
上海理工大学 Sci. Adv.｜自旋解耦超表面赋能 D-D²NN：大容量方向依赖全光衍射神经网络
人工智能·深度学习·神经网络·超表面
AI棒棒牛4 小时前
基于 IDEAL 思想的 YOLO26 双分支适配：差异证据与不确定性门控，RGB+IR双分支创新！
人工智能·深度学习·yolo·目标检测·yolo26
WenGyyyL4 小时前
从零理解强化学习：为什么 AI 需要 Advantage，而不只是 Reward？
人工智能·python·深度学习·机器学习·nlp
热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02如何新建文件夹？ 电脑新建文件夹的4种方法03GitHub 镜像站点04历年考研数学一、数学二、数学三真题试卷及答案PDF05几个好用的ip纯净度检测网站06ai安全工具:CyberStrikeAI安装部署与使用07苹果iPA游戏软件资源下载网站08昇腾910B系列芯片完全指南：四款型号对比、服务器选型与大模型部署实战09「全语录+省流版」邵艾伦 × 孙宇晨：聊风口、聊个人IP、聊年轻人如何抓住 AI 时代的机会10乙二醇水溶液物性参数