重要更新

2022年2月24日，俄罗斯宣布开战并入侵和平的乌克兰。继吞并克里米亚和占领顿巴斯地区之后，普京政权决定摧毁乌克兰民族认同。乌克兰人民展现出顽强的抵抗精神，向全世界展示了为国家独立而战的意义。

乌克兰政府开通了一个网站，帮助俄罗斯的母亲、妻子和姐妹寻找在乌克兰阵亡或被俘的亲人——https://200rf.com 和 https://t.me/rf200_now（Telegram频道）。我们的目标是让仍身处俄罗斯和白俄罗斯的人们了解真相，从而拒绝参与对乌克兰的侵略。

请帮助我们让更多人了解乌克兰正在发生的事情——“俄罗斯世界”给乌克兰带来的暴力与非人道恐怖行径。这里有一个全面的维基页面，介绍你可以如何帮助结束这场战争：https://how-to-help-ukraine-now.super.site/

官方渠道：

• 乌克兰议会官方账号
• 乌克兰国防部
• 乌克兰总统办公室
• 乌克兰内阁
• 战略传播中心
• 乌克兰外交部长

荣耀归于乌克兰！

Pytorch-toolbelt

pytorch-toolbelt 是一个基于 PyTorch 的 Python 库，提供了一系列实用工具和模块，旨在加速研发原型设计和 Kaggle 竞赛中的模型训练：

包含内容

• 基于灵活编码器-解码器架构的便捷模型构建。
• 模块：CoordConv、SCSE、Hypercolumn、深度可分离卷积等。
• 针对分割和分类任务的 GPU 友好型测试时增强 TTA。
• 对超大尺寸图像（如 5000x5000）进行 GPU 加速推理。
• 常用工具函数：固定/恢复随机种子、文件系统工具、评估指标等。
• 损失函数：BinaryFocalLoss、Focal、ReducedFocal、Lovasz、Jaccard 和 Dice 损失、Wing Loss 等。
• 与 Catalyst 库的集成扩展功能（批量预测可视化、额外指标）。

示例：Catalyst、Albumentations、Pytorch Toolbelt 示例：CamVid 数据集上的语义分割

缘由

坦率地说，“我需要一种方便的方式来复用代码，以支持我的 Kaggle 职业生涯”。2018 年，我获得了 Kaggle Master 称号，这是一段漫长的过程。我经常发现自己一遍又一遍地重复使用旧的代码流程。最终，这些经验积累形成了这个库。

该库并非旨在取代 Catalyst、Ignite 或 Fast.ai 等高级框架，而是作为它们的补充。

安装

pip install pytorch_toolbelt

如何...

模型创建

创建 U-Net 编码器-解码器模型

以下代码片段创建了一个用于二分类分割的原生 U-Net 模型。按照设计，编码器和解码器都会输出一系列特征图，从精细（高分辨率，索引为 0）到粗糙（低分辨率）。访问所有中间特征图非常有用，例如在应用深度监督损失时，或者在目标检测任务中，中间特征图的访问往往是必要的。

from torch import nn
from pytorch_toolbelt.modules import encoders as E
from pytorch_toolbelt.modules import decoders as D

class UNet(nn.Module):
    def __init__(self, input_channels, num_classes):
        super().__init__()
        self.encoder = E.UnetEncoder(in_channels=input_channels, out_channels=32, growth_factor=2)
        self.decoder = D.UNetDecoder(self.encoder.channels, decoder_features=32)
        self.logits = nn.Conv2d(self.decoder.channels[0], num_classes, kernel_size=1)

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return self.logits(x[0])

创建带有预训练编码器的 FPN 编码器-解码器模型

与前一个示例类似，你可以将解码器替换为带有特征融合的 FPN 解码器。

from torch import nn
from pytorch_toolbelt.modules import encoders as E
from pytorch_toolbelt.modules import decoders as D

class SEResNeXt50FPN(nn.Module):
    def __init__(self, num_classes, fpn_channels):
        super().__init__()
        self.encoder = E.SEResNeXt50Encoder()
        self.decoder = D.FPNCatDecoder(self.encoder.channels, fpn_channels)
        self.logits = nn.Conv2d(self.decoder.channels[0], num_classes, kernel_size=1)

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return self.logits(x[0])

更改编码器的输入通道数

pytorch_toolbelt 中的所有编码器都支持更改输入通道数。只需调用 encoder.change_input_channels(num_channels)，即可修改第一层卷积的输入通道数。在可能的情况下，现有卷积层的权重会被重用（如果新通道数大于默认值，则会用随机初始化的权重填充新增部分）。该方法返回 self，因此可以链式调用。

from pytorch_toolbelt.modules import encoders as E

encoder = E.SEResnet101Encoder()
encoder = encoder.change_input_channels(6)

其他

统计编码器/解码器及其他模块的参数量

在设计模型并优化神经网络中的特征数量时，我发现打印高层模块（如编码器和解码器）的参数量非常有帮助。以下是使用 pytorch_toolbelt 实现的方法：

from torch import nn
from pytorch_toolbelt.modules import encoders as E
from pytorch_toolbelt.modules import decoders as D
from pytorch_toolbelt.utils import count_parameters

class SEResNeXt50FPN(nn.Module):
    def __init__(self, num_classes, fpn_channels):
        super().__init__()
        self.encoder = E.SEResNeXt50Encoder()
        self.decoder = D.FPNCatDecoder(self.encoder.channels, fpn_channels)
        self.logits = nn.Conv2d(self.decoder.channels[0], num_classes, kernel_size=1)

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return self.logits(x[0])

net = SEResNeXt50FPN(1, 128)
print(count_parameters(net))
# 输出：{'total': 34232561, 'trainable': 34232561, 'encoder': 25510896, 'decoder': 8721536, 'logits': 129}

组合多个损失函数

虽然像 Catalyst 这样的高级深度学习框架提供了更灵活的方式来组合多损失，但这里是一个纯 PyTorch 实现的多损失组合方法：

from pytorch_toolbelt import losses as L

# 创建一个损失函数，它是焦点损失和 Lovasz 损失的加权和，权重分别为 1.0 和 0.5。
loss = L.JointLoss(L.FocalLoss(), L.LovaszLoss(), 1.0, 0.5)

TTA / 推理

为模型应用测试时增强（TTA）

测试时增强（TTA）可以在训练和测试阶段使用。

from pytorch_toolbelt.inference import tta

model = UNet()

# 使用水平翻转实现真正有效的图像分类 TTA：
logits = tta.fliplr_image2label(model, input)

# 使用 D4 增强实现真正有效的图像分割 TTA：
logits = tta.d4_image2mask(model, input)

对超大图像进行推理：

在许多情况下，需要对非常大的图像（5000像素及以上）进行分割。处理如此大的像素数组会遇到几个问题：

1. CUDA 张量的最大尺寸存在限制（具体数值取决于驱动程序和 GPU 版本）。
2. 复杂的 CNN 架构在推理相对较小的 1024x1024 图像时，可能会轻易耗尽所有可用的 GPU 内存，从而无法处理更大分辨率的图像。

一种解决方案是将输入图像切分成小块（可选重叠），分别送入模型进行推理，然后将结果拼接起来。这样既可以确保 GPU 内存的使用上限，又能够处理任意大小的图像。

import numpy as np
from torch.utils.data import DataLoader
import cv2

from pytorch_toolbelt.inference.tiles import ImageSlicer, CudaTileMerger
from pytorch_toolbelt.utils.torch_utils import tensor_from_rgb_image, to_numpy


image = cv2.imread('really_huge_image.jpg')
model = get_model(...)

# 将大图像切割成重叠的小块
tiler = ImageSlicer(image.shape, tile_size=(512, 512), tile_step=(256, 256))

# HCW -> CHW。可选在此处进行归一化
tiles = [tensor_from_rgb_image(tile) for tile in tiler.split(image)]

# 分配一个 CUDA 缓冲区来存储完整的掩码
merger = CudaTileMerger(tiler.target_shape, 1, tiler.weight)

# 对每个小块进行预测并累积结果
for tiles_batch, coords_batch in DataLoader(list(zip(tiles, tiler.crops)), batch_size=8, pin_memory=True):
    tiles_batch = tiles_batch.float().cuda()
    pred_batch = model(tiles_batch)

    merger.integrate_batch(pred_batch, coords_batch)

# 对累积的掩码进行归一化，并转换回 NumPy 数组
merged_mask = np.moveaxis(to_numpy(merger.merge()), 0, -1).astype(np.uint8)
merged_mask = tiler.crop_to_orignal_size(merged_mask)

高级示例

1. Inria 卫星图像分割
2. CamVid 语义分割

引用

@misc{Khvedchenya_Eugene_2019_PyTorch_Toolbelt,
  author = {Khvedchenya, Eugene},
  title = {PyTorch Toolbelt},
  year = {2019},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/BloodAxe/pytorch-toolbelt}},
  commit = {cc5e9973cdb0dcbf1c6b6e1401bf44b9c69e13f3}
}

PyTorch Toolbelt 快速上手指南

pytorch-toolbelt 是一个专为 PyTorch 设计的工具库，旨在加速研发原型设计和 Kaggle 竞赛。它提供了灵活的编码器 - 解码器架构、丰富的模块（如 CoordConv, SCSE）、高效的测试时增强（TTA）以及多种常用的损失函数。

环境准备

• 系统要求：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.7+
• 前置依赖：
  • PyTorch (建议最新版本)
  • torchvision
  • NumPy
  • OpenCV-Python (cv2) (用于大图像推理示例)

请确保已正确安装 CUDA 驱动及对应的 PyTorch GPU 版本以获得最佳性能。

安装步骤

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

通用安装命令：

pip install pytorch-toolbelt

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

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

基本使用

1. 构建编码器 - 解码器模型 (U-Net)

pytorch-toolbelt 的核心优势在于快速搭建基于 Encoder-Decoder 架构的模型。以下是一个创建用于二分类分割的 vanilla U-Net 的示例：

from torch import nn
from pytorch_toolbelt.modules import encoders as E
from pytorch_toolbelt.modules import decoders as D

class UNet(nn.Module):
    def __init__(self, input_channels, num_classes):
        super().__init__()
        # 定义编码器：输入通道数，基础输出通道数，增长因子
        self.encoder = E.UnetEncoder(in_channels=input_channels, out_channels=32, growth_factor=2)
        # 定义解码器：使用编码器的通道配置
        self.decoder = D.UNetDecoder(self.encoder.channels, decoder_features=32)
        # 定义最终输出层
        self.logits = nn.Conv2d(self.decoder.channels[0], num_classes, kernel_size=1)

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        # 取最高分辨率的特征图进行预测
        return self.logits(x[0])

# 实例化模型
model = UNet(input_channels=3, num_classes=1)

2. 组合损失函数

库中内置了多种常用的损失函数（如 Focal Loss, Lovasz Loss, Dice Loss 等），并支持轻松组合：

from pytorch_toolbelt import losses as L

# 创建一个加权组合损失：Focal Loss (权重 1.0) + Lovasz Loss (权重 0.5)
loss_fn = L.JointLoss(L.FocalLoss(), L.LovaszLoss(), 1.0, 0.5)

# 在训练循环中使用
# output = model(input)
# loss = loss_fn(output, target)

3. 统计模型参数量

在设计模型时，快速查看各部分（Encoder/Decoder）的参数量非常有用：

from pytorch_toolbelt.utils import count_parameters

# 假设已经定义了上述 model
params_info = count_parameters(model)
print(params_info)
# 输出示例：{'total': ..., 'trainable': ..., 'encoder': ..., 'decoder': ...}

4. 测试时增强 (TTA)

利用内置的 TTA 功能提升推理效果，支持分类和分割任务：

from pytorch_toolbelt.inference import tta

# 图像分类：使用水平翻转 TTA
# logits = tta.fliplr_image2label(model, input_tensor)

# 图像分割：使用 D4 (旋转 + 翻转) 增强 TTA
# mask_logits = tta.d4_image2mask(model, input_tensor)

5. 超大图像推理

针对无法一次性放入显存的超大图像（如 5000x5000 像素），库提供了切片推理与合并工具：

import numpy as np
import cv2
from torch.utils.data import DataLoader
from pytorch_toolbelt.inference.tiles import ImageSlicer, CudaTileMerger
from pytorch_toolbelt.utils.torch_utils import tensor_from_rgb_image, to_numpy

# 读取大图
image = cv2.imread('really_huge_image.jpg')
# model = get_model(...) 

# 1. 将大图切割为重叠的瓦片 (tile)
tiler = ImageSlicer(image.shape, tile_size=(512, 512), tile_step=(256, 256))
tiles = [tensor_from_rgb_image(tile) for tile in tiler.split(image)]

# 2. 初始化 CUDA 合并器
merger = CudaTileMerger(tiler.target_shape, num_classes=1, weight=tiler.weight)

# 3. 分批推理并合并结果
for tiles_batch, coords_batch in DataLoader(list(zip(tiles, tiler.crops)), batch_size=8, pin_memory=True):
    tiles_batch = tiles_batch.float().cuda()
    with torch.no_grad():
        pred_batch = model(tiles_batch)
    merger.integrate_batch(pred_batch, coords_batch)

# 4. 获取最终合并后的掩码
merged_mask = np.moveaxis(to_numpy(merger.merge()), 0, -1).astype(np.uint8)
merged_mask = tiler.crop_to_orignal_size(merged_mask)