PyTorch 实战：图像分割从入门到部署

图像分割要求模型对每个像素进行分类，是计算机视觉中最精细的任务之一。U-Net 以其编码器-解码器结构和跳跃连接，成为医学图像和遥感图像分割的首选架构。

U-Net 的核心实现：

class UNet(nn.Module):
    def __init__(self, in_channels=3, num_classes=2):
        super().__init__()
        # 编码器
        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        # 解码器
        self.dec3 = self.up_block(256, 128)
        self.dec2 = self.up_block(128, 64)
        self.final = nn.Conv2d(64, num_classes, 1)

    def conv_block(self, in_c, out_c):
        return nn.Sequential(
            nn.Conv2d(in_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.ReLU(inplace=True),
        )

数据增强对分割任务至关重要。我们使用 Albumentations 库实现同步的图像和掩码增强：

transform = A.Compose([
    A.RandomCrop(256, 256),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.3),
    A.ElasticTransform(p=0.3),
    A.Normalize(),
])

模型部署方面，我们通过 ONNX 导出和 TensorRT 优化，将推理延迟从 45ms 降低到 8ms，满足了实时分割的需求。关键步骤包括：固定 BatchNorm、融合卷积层、选择 FP16 精度模式。