pytorch-toolbelt

PyTorch-Toolbelt 项目介绍

PyTorch-Toolbelt 是一个专为 PyTorch 设计的 Python 库，它带来了一系列便捷的工具和模块，帮助用户加速研究与开发过程，特别是在 Kaggle 数据竞赛中显得尤为实用。

项目亮点

模型构建

PyTorch-Toolbelt 通过灵活的编码器-解码器架构，简化了模型的构建。这种架构通常用于图像分割和其它需要逐步处理高分辨率特征图的任务。

模块支持

库中包含了各种实用模块，例如 CoordConv、SCSE、Hypercolumn，以及深度可分离卷积等。这些模块可以帮助用户在模型中实现更加复杂和先进的特征处理。

GPU 友好的增强与推理

工具包支持分割和分类任务的测试时增强（Test-Time Augmentation, TTA），还有针对巨大图像（如5000x5000像素）的推理优化，这使得在大规模图像上进行机器学习预测成为可能。

常用函数与损失函数

PyTorch-Toolbelt 提供了一系列常用的函数和损失函数，包括二元焦点损失（BinaryFocalLoss）、Jaccard 和 Dice 损失、Wing 损失等。这些工具可以帮助用户更好地处理数据、优化模型表现。

Catalyst 库的扩展

它为功能强大的机器学习训练库 Catalyst 提供了一些扩展，例如批次预测的可视化和附加的评估指标。

为什么创建 PyTorch-Toolbelt

工具包的创建初衷是为了解决作者自身在 Kaggle 生涯中代码复用时遇到的问题。作者通过不断地复用和优化旧代码，逐渐形成了这一仓库。值得注意的是，PyTorch-Toolbelt 并非要取代现有的高层框架如 Catalyst、Ignite 或 Fast.ai，它的目标是作为它们的补充，提供更便捷的工具和模块。

安装

用户可以通过简单的命令来安装 PyTorch-Toolbelt：

pip install pytorch_toolbelt

使用范例

创建 Encoder-Decoder U-Net 模型

以下是一个用于二元分割任务的基础 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])

大图推理

在处理5000像素以上的巨大图像时，可以将图像切片成小块，然后逐块进行推理，最后合并这些结果：

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))
tiles = [tensor_from_rgb_image(tile) for tile in tiler.split(image)]
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)

merged_mask = np.moveaxis(to_numpy(merger.merge()), 0, -1).astype(np.uint8)
merged_mask = tiler.crop_to_orignal_size(merged_mask)