TTAch

使用 PyTorch 进行测试时数据增强！

与数据增强对训练集的作用类似，测试时数据增强的目的是对测试图像进行随机变换。因此，我们不会只向训练好的模型展示一次常规的“干净”图像，而是会多次展示经过增强后的图像。然后，我们将对每张对应图像的预测结果取平均值，并将其作为最终的预测结果 [1]。

           输入
             |           # 输入图像批次 
        / / /|\ \ \      # 应用增强操作（翻转、旋转、缩放等）
       | | | | | | |     # 将增强后的批次输入模型
       | | | | | | |     # 对每个批次的掩码/标签进行逆变换
        \ \ \ / / /      # 合并预测结果（均值、最大值、几何平均等）
             |           # 输出掩码/标签批次
           输出

目录

1. 快速入门
2. 变换
3. 别名
4. 合并模式
5. 安装

快速入门

分割模型封装 [文档字符串]:

import ttach as tta
tta_model = tta.SegmentationTTAWrapper(model, tta.aliases.d4_transform(), merge_mode='mean')

分类模型封装 [文档字符串]:

tta_model = tta.ClassificationTTAWrapper(model, tta.aliases.five_crop_transform())

关键点模型封装 [文档字符串]:

tta_model = tta.KeypointsTTAWrapper(model, tta.aliases.flip_transform(), scaled=True)

注意: 模型必须以 torch([x1, y1, ..., xn, yn]) 的格式返回关键点。

高级示例

自定义变换:

# 定义了 2 * 2 * 3 * 3 = 36 种增强方式！
transforms = tta.Compose(
    [
        tta.HorizontalFlip(),
        tta.Rotate90(angles=[0, 180]),
        tta.Scale(scales=[1, 2, 4]),
        tta.Multiply(factors=[0.9, 1, 1.1]),        
    ]
)

tta_model = tta.SegmentationTTAWrapper(model, transforms)

自定义模型（多输入/多输出）

# 示例：如何使用 TTA 处理一批图像
# 其中 `image`/`mask` 是 4D 张量 (B, C, H, W)，`label` 是 2D 张量 (B, N)

for transformer in transforms: # 自定义变换或例如 tta.aliases.d4_transform() 
    
    # 增强图像
    augmented_image = transformer.augment_image(image)
    
    # 输入模型
    model_output = model(augmented_image, another_input_data)
    
    # 对掩码和标签进行逆增强
    deaug_mask = transformer.deaugment_mask(model_output['mask'])
    deaug_label = transformer.deaugment_label(model_output['label'])
    
    # 保存结果
    labels.append(deaug_mask)
    masks.append(deaug_label)
    
# 按照需求对结果进行聚合，例如取均值/最大值/最小值
label = mean(labels)
mask = mean(masks)

变换

变换	参数	取值
HorizontalFlip	-	-
VerticalFlip	-	-
Rotate90	角度	列表[0, 90, 180, 270]
Scale	缩放比例 插值方法	列表[浮点数] "nearest"/"linear"
Resize	目标尺寸 原始尺寸 插值方法	列表[元组[int, int]] 元组[int,int] "nearest"/"linear"
Add	值	列表[浮点数]
Multiply	因子	列表[浮点数]
FiveCrops	裁剪高度 裁剪宽度	整数 整数

别名

• flip_transform（水平+垂直翻转）
• hflip_transform（水平翻转）
• d4_transform（翻转+旋转 0°、90°、180°、270°）
• multiscale_transform（尺度变换，可指定缩放比例）
• five_crop_transform（四个角裁剪+中心裁剪）
• ten_crop_transform（五个角裁剪+水平翻转后的五个角裁剪）

合并模式

• mean
• gmean（几何平均）
• sum
• max
• min
• tsharpen（温度锐化 方法，t=0.5）

安装

PyPI:

$ pip install ttach

源码:

$ pip install git+https://github.com/qubvel/ttach

运行测试

docker build -f Dockerfile.dev -t ttach:dev . && docker run --rm ttach:dev pytest -p no:cacheprovider

ttach 快速上手指南

ttach 是一个基于 PyTorch 的图像测试时增强（Test Time Augmentation, TTA）库。它通过对测试图像进行随机变换（如翻转、旋转、缩放等），将增强后的多批次图像输入模型，最后合并预测结果以提升模型性能。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：建议 Python 3.7+
• 核心依赖：
  • PyTorch (>= 1.0)
  • torchvision
• 前置要求：确保已安装好 PyTorch 深度学习环境，且你的模型能够接受标准的 Tensor 输入并输出预测结果。

安装步骤

推荐使用 pip 进行安装。国内用户可使用清华或阿里镜像源加速下载。

使用官方源安装：

pip install ttach

使用国内镜像源安装（推荐）：

pip install ttach -i https://pypi.tuna.tsinghua.edu.cn/simple

从源码安装（获取最新开发版）：

pip install git+https://github.com/qubvel/ttach

基本使用

ttach 提供了针对不同任务的封装器（Wrapper），只需几行代码即可将现有模型升级为支持 TTA 的模型。

1. 语义分割模型 (Segmentation)

使用预定义的 d4_transform（包含翻转和 90 度旋转）并将预测结果取平均值：

import ttach as tta

# model 为你已经训练好的 PyTorch 分割模型
tta_model = tta.SegmentationTTAWrapper(model, tta.aliases.d4_transform(), merge_mode='mean')

# 推理过程与普通模型一致
with torch.no_grad():
    prediction = tta_model(input_image)

2. 分类模型 (Classification)

使用 five_crop_transform（五裁剪增强）：

import ttach as tta

# model 为你已经训练好的 PyTorch 分类模型
tta_model = tta.ClassificationTTAWrapper(model, tta.aliases.five_crop_transform())

with torch.no_grad():
    prediction = tta_model(input_image)

3. 关键点检测模型 (Keypoints)

使用水平翻转增强，注意模型需返回格式为 torch([x1, y1, ..., xn, yn]) 的关键点：

import ttach as tta

# scaled=True 表示关键点坐标会随图像缩放自动调整
tta_model = tta.KeypointsTTAWrapper(model, tta.aliases.flip_transform(), scaled=True)

with torch.no_grad():
    prediction = tta_model(input_image)

进阶：自定义增强组合

你可以自由组合多种变换策略。以下示例构建了一个包含 36 种增强组合的策略：

import ttach as tta

transforms = tta.Compose(
    [
        tta.HorizontalFlip(),
        tta.Rotate90(angles=[0, 180]),
        tta.Scale(scales=[1, 2, 4]),
        tta.Multiply(factors=[0.9, 1, 1.1]),        
    ]
)

tta_model = tta.SegmentationTTAWrapper(model, transforms, merge_mode='mean')

常用合并模式 (merge_mode)：

• mean: 算术平均（默认）
• gmean: 几何平均
• max: 取最大值
• min: 取最小值
• tsharpen: 温度锐化