从零构建Deeplabv3+ ResNet50：一份面向实战的图像分割全流程指南

如果你正在寻找一份能让你真正动手跑起来，并且能深刻理解其内部运作机制的Deeplabv3+ ResNet50实现指南，那么你来对地方了。这篇文章不会仅仅停留在理论复述上，而是会带你从环境搭建开始，一步步构建、训练并调试一个完整的图像分割模型。我遇到过太多初学者，他们能看懂论文里的公式，却卡在代码调试和实际数据适配上。因此，这里的内容会聚焦于那些在官方文档和论文里找不到的“坑”，以及如何用最清晰的代码结构来理解这个强大的语义分割架构。无论你是想快速在自己的数据集上验证想法，还是希望深入理解空洞卷积和ASPP模块的细节，这篇文章都将提供一条直达目标的路径。

1. 环境准备与项目初始化

在开始敲代码之前，一个稳定、可复现的环境是高效工作的基石。我强烈建议使用 conda 或 venv 来管理你的Python环境，这能避免不同项目间的依赖冲突。对于PyTorch，直接访问其官方网站获取安装命令是最稳妥的方式，记得根据你的CUDA版本进行选择。

注意：如果你的GPU显存有限（例如小于8GB），在后续训练时可能需要调小batch_size，或者考虑使用梯度累积技术。

除了PyTorch，我们还需要一些辅助工具库。下面是一个推荐的基础环境配置清单：

# 创建并激活conda环境（示例）
conda create -n deeplabv3 python=3.8
conda activate deeplabv3

# 安装PyTorch（请根据你的CUDA版本调整，此处以CUDA 11.3为例）
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113

# 安装其他必要的库
pip install opencv-python pillow matplotlib tqdm tensorboard scikit-learn pandas

项目目录结构也值得在开始前规划好。一个清晰的结构不仅能让你自己思路清晰，也方便与他人协作。我通常采用如下结构：

deeplabv3_project/
├── configs/               # 配置文件，用于管理超参数
│   └── train_config.yaml
├── data/                  # 数据相关
│   ├── train_images/      # 训练集原图
│   ├── train_masks/       # 训练集标注图
│   ├── val_images/        # 验证集原图
│   └── val_masks/         # 验证集标注图
├── models/                # 模型定义
│   ├── __init__.py
│   ├── backbone.py        # ResNet50 Backbone
│   ├── aspp.py            # ASPP模块
│   └── deeplabv3.py       # 完整的Deeplabv3+模型
├── utils/                 # 工具函数
│   ├── dataset.py         # 自定义Dataset类
│   ├── transforms.py      # 数据增强
│   └── metrics.py         # 评估指标计算
├── scripts/               # 执行脚本
│   ├── train.py
│   └── evaluate.py
├── outputs/               # 训练输出（日志、模型权重、可视化结果）
│   ├── logs/
│   ├── checkpoints/
│   └── predictions/
└── requirements.txt

这种模块化的设计，使得数据加载、模型定义、训练逻辑相互解耦，后期调试和扩展都会轻松很多。

2. 核心模块深度解析与自实现

很多教程直接调用 torchvision.models.segmentation.deeplabv3_resnet50，这固然方便，但如同一个黑盒，出了问题难以排查。要真正掌握它，我们必须亲手搭建其核心部件。Deeplabv3+ 的精髓主要在于空洞卷积（Dilated/Atrous Convolution）和空洞空间金字塔池化（ASPP）模块，它们被巧妙地集成在ResNet50骨干网络上。

2.1 改造ResNet50作为Backbone

标准的ResNet50是为图像分类设计的，其最后两个阶段（layer3和layer4）会进行下采样，导致特征图空间分辨率损失严重，这对于需要精细位置信息的语义分割是致命的。Deeplabv3的解决方案是引入空洞卷积。

首先，我们需要修改ResNet50的最后两个阶段，将普通卷积替换为空洞卷积，并移除不必要的下采样。具体来说：

• Layer3: 将其中的3x3卷积的dilation参数设置为2，padding也相应调整为2，以保持特征图尺寸不变。
• Layer4: 将其中的3x3卷积的dilation参数设置为4，padding调整为4。同时，要移除该阶段第一个Bottleneck中的下采样（将stride=2改为stride=1）。

下面是一个关键代码片段，展示了如何修改PyTorch自带的ResNet来获取我们需要的Backbone：

import torch
import torch.nn as nn
from torchvision.models import resnet50

def make_dilated_resnet50(pretrained=True):
    model = resnet50(pretrained=pretrained)
    # 获取layer3和layer4
    layer3 = model.layer3
    layer4 = model.layer4

    # 修改layer3: 将所有3x3卷积的dilation和padding改为2
    for bottleneck in layer3:
        for conv in [bottleneck.conv2]:
            if conv.kernel_size == (3, 3):
                conv.dilation = (2, 2)
                conv.padding = (2, 2)

    # 修改layer4: 首先移除第一个bottleneck的下采样
    # 第一个bottleneck的downsample层和conv2的stride需要修改
    first_bottleneck = layer4[0]
    if first_bottleneck.downsample is not None:
        # 将下采样卷积的stride从2改为1
        first_bottleneck.downsample[0].stride = (1, 1)
    # 将第一个bottleneck中conv2的stride从2改为1
    first_bottleneck.conv2.stride = (1, 1)

    # 然后修改layer4所有bottleneck中3x3卷积的dilation和padding为4
    for bottleneck in layer4:
        for conv in [bottleneck.conv2]:
            if conv.kernel_size == (3, 3):
                conv.dilation = (4, 4)
                conv.padding = (4, 4)

    # 我们只需要特征提取部分，去掉最后的全局平均池化和全连接层
    backbone = nn.Sequential(
        model.conv1,
        model.bn1,
        model.relu,
        model.maxpool,
        model.layer1,
        model.layer2,
        layer3,  # 使用修改后的layer3
        layer4   # 使用修改后的layer4
    )
    return backbone

# 测试Backbone输出
backbone = make_dilated_resnet50(pretrained=True)
dummy_input = torch.randn(2, 3, 512, 512)
with torch.no_grad():
    features = backbone(dummy_input)
print(f"Backbone输出特征图形状: {features.shape}")  # 期望输出: torch.Size([2, 2048, 32, 32])

通过这样的改造，输入图像经过Backbone后，其空间尺寸仅下采样了16倍（例如512x512 -> 32x32），而不是原始ResNet50的32倍，保留了更多空间信息。

2.2 构建ASPP模块：捕捉多尺度上下文

ASPP模块是Deeplabv3系列的核心创新。它的思想是并行的使用多个具有不同采样率的空洞卷积层以及一个全局平均池化层，来捕获图像中不同尺度的上下文信息。这些并行的分支输出会被拼接起来，再通过一个1x1卷积进行融合。

一个典型的ASPP模块包含以下五个分支：

1. 一个1x1卷积（空洞率=1，即普通卷积）。
2. 一个3x3空洞卷积，空洞率=12。
3. 一个3x3空洞卷积，空洞率=24。
4. 一个3x3空洞卷积，空洞率=36。
5. 一个全局平均池化层，后接1x1卷积和双线性上采样。

不同空洞率对应着不同的感受野，空洞率越大，卷积核“看到”的图像范围就越广，从而能捕捉更全局的上下文。而全局平均池化分支则提供了图像级别的全局上下文信息。

以下是ASPP模块的PyTorch实现：

import torch.nn.functional as F

class ASPP(nn.Module):
    def __init__(self, in_channels, out_channels=256, rates=[6, 12, 18]):
        super(ASPP, self).__init__()
        # 分支1: 1x1卷积
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        # 分支2-4: 不同空洞率的3x3卷积
        self.convs = nn.ModuleList()
        for rate in rates:
            self.convs.append(
                nn.Sequential(
                    nn.Conv2d(in_channels, out_channels, 3, padding=rate, dilation=rate, bias=False),
                    nn.BatchNorm2d(out_channels),
                    nn.ReLU(inplace=True)
                )
            )
        # 分支5: 全局特征
        self.global_avg_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1), # 输出 (B, C, 1, 1)
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        # 融合所有分支特征的1x1卷积
        self.project = nn.Sequential(
            nn.Conv2d(out_channels * (len(rates) + 2), out_channels, 1, bias=False), # +2 是conv1和global分支
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5)
        )

    def forward(self, x):
        h, w = x.shape[2:]
        # 处理前四个卷积分支
        out = [self.conv1(x)]
        for conv in self.convs:
            out.append(conv(x))
        # 处理全局池化分支并上采样
        global_feat = self.global_avg_pool(x)
        global_feat = F.interpolate(global_feat, size=(h, w), mode='bilinear', align_corners=False)
        out.append(global_feat)
        # 沿通道维度拼接
        out = torch.cat(out, dim=1)
        # 投影融合
        out = self.project(out)
        return out

# 测试ASPP模块
aspp = ASPP(in_channels=2048, out_channels=256, rates=[6, 12, 18])
dummy_backbone_out = torch.randn(2, 2048, 32, 32) # 假设Backbone输出
aspp_out = aspp(dummy_backbone_out)
print(f"ASPP输出特征图形状: {aspp_out.shape}")  # 期望输出: torch.Size([2, 256, 32, 32])

这里我使用了[6, 12, 18]的空洞率，这是一个在多个数据集上表现良好的经验值。你可以根据自己任务中目标物体的大小调整这些值。

2.3 组装完整的Deeplabv3+模型

有了改造后的Backbone和ASPP模块，我们就可以组装完整的Deeplabv3+模型了。Deeplabv3+相比v3，增加了一个来自骨干网络中间层的低级特征，与经过ASPP处理后的高级语义特征进行融合，以恢复更精细的边缘信息。

通常，我们从ResNet的layer1（或layer2）的输出中提取低级特征。完整的模型构建流程如下：

1. 图像输入经过修改后的ResNet50 Backbone。
2. 从Backbone的中间层（如layer1）提取低级特征。
3. Backbone的最终输出送入ASPP模块，得到高级语义特征。
4. 将高级语义特征上采样到与低级特征相同的尺寸。
5. 将上采样后的高级特征与低级特征在通道维度拼接。
6. 通过几个卷积层进行融合和最终预测。

class DeepLabV3Plus(nn.Module):
    def __init__(self, num_classes, backbone='resnet50', pretrained=True):
        super(DeepLabV3Plus, self).__init__()
        # 1. 构建Backbone并获取中间层
        self.backbone = make_dilated_resnet50(pretrained=pretrained)
        # 假设我们取layer1的输出作为低级特征 (通道数256)
        self.low_level_features = None # 我们会在forward中捕获它

        # 2. ASPP模块
        self.aspp = ASPP(in_channels=2048, out_channels=256)

        # 3. 解码器部分
        # 首先处理低级特征，降低其通道数
        self.low_level_conv = nn.Sequential(
            nn.Conv2d(256, 48, 1, bias=False),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True)
        )
        # 融合高低级特征的卷积块
        self.decoder_conv = nn.Sequential(
            nn.Conv2d(256+48, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1)
        )
        # 最终的分类卷积层
        self.classifier = nn.Conv2d(256, num_classes, 1)

    def forward(self, x):
        h, w = x.shape[2:]
        # 提取特征 - 这里需要手动获取中间层输出
        # 为了清晰，我们模拟一个过程：实际实现时可能需要修改backbone的forward函数以返回多级特征
        # 此处为简化示意，假设self.backbone能返回(low_level_feat, high_level_feat)
        low_level_feat, aspp_input = self._extract_features(x) # 假设的辅助函数

        # 处理高级特征
        aspp_out = self.aspp(aspp_input) # [B, 256, H/16, W/16]

        # 上采样ASPP输出到低级特征的尺寸
        aspp_out_up = F.interpolate(aspp_out, size=low_level_feat.shape[2:], mode='bilinear', align_corners=False)

        # 处理低级特征
        low_level_feat = self.low_level_conv(low_level_feat) # [B, 48, H/4, W/4]

        # 拼接特征
        fused = torch.cat([aspp_out_up, low_level_feat], dim=1) # [B, 304, H/4, W/4]
        fused = self.decoder_conv(fused) # [B, 256, H/4, W/4]

        # 最终预测并上采样回原图尺寸
        out = self.classifier(fused) # [B, num_classes, H/4, W/4]
        out = F.interpolate(out, size=(h, w), mode='bilinear', align_corners=False)

        return out

    def _extract_features(self, x):
        # 这是一个示意函数。在实际完整实现中，你需要重写Backbone的forward，
        # 使其能返回中间层和最终层的特征图。
        # 例如，可以依次执行backbone的各个层并保存中间结果。
        pass

提示：在实际编码时，更优雅的做法是创建一个新的Backbone类，继承自修改后的ResNet，并重写其forward方法，使其返回一个包含多级特征的字典或元组。这比在完整模型里拆解网络层要清晰得多。

3. 数据管道与训练策略实战

模型搭建好了，但要让它在你的数据上work起来，数据准备和训练策略同样关键。语义分割任务对数据标注的要求很高，通常需要像素级的掩码图。

3.1 构建高效的数据加载器

我们使用PyTorch的Dataset和DataLoader。数据增强对于分割任务至关重要，它能有效防止过拟合并提升模型泛化能力。对于输入图像和对应的掩码，必须施加完全相同的空间变换。

from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
import torchvision.transforms.functional as TF
import random

class SegmentationDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.images = sorted([f for f in os.listdir(image_dir) if f.endswith('.jpg') or f.endswith('.png')])
        self.masks = sorted([f for f in os.listdir(mask_dir) if f.endswith('.png')])

        # 简单检查文件名是否对应
        assert len(self.images) == len(self.masks), "图像和掩码数量不匹配"
        for img, msk in zip(self.images, self.masks):
            # 假设文件名（不含后缀）相同
            if os.path.splitext(img)[0] != os.path.splitext(msk)[0]:
                print(f"警告: 可能不匹配的图像-掩码对: {img} vs {msk}")

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.images[idx])
        mask_path = os.path.join(self.mask_dir, self.masks[idx])

        image = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L") # 灰度图，单通道

        # 基础转换：转为Tensor
        image = TF.to_tensor(image)
        # 对于掩码，我们通常需要将其像素值转换为类别索引
        # 假设掩码是单通道，背景为0，不同物体为1,2,3...
        mask = torch.from_numpy(np.array(mask)).long()

        if self.transform:
            # 对于需要同步变换的数据增强（如旋转、翻转），需要自定义transform函数
            image, mask = self._apply_transform(image, mask)

        return image, mask

    def _apply_transform(self, image, mask):
        """应用需要同步的数据增强"""
        # 随机水平翻转
        if random.random() > 0.5:
            image = TF.hflip(image)
            mask = TF.hflip(mask.unsqueeze(0)).squeeze(0) # 为mask增加再移除通道维度以使用TF

        # 随机垂直翻转
        if random.random() > 0.5:
            image = TF.vflip(image)
            mask = TF.vflip(mask.unsqueeze(0)).squeeze(0)

        # 随机旋转 (-10度到10度之间)
        angle = random.uniform(-10, 10)
        image = TF.rotate(image, angle)
        mask = TF.rotate(mask.unsqueeze(0), angle).squeeze(0)

        return image, mask

对于更复杂的数据增强（如随机缩放、裁剪、颜色抖动），可以考虑使用albumentations库，它专门为图像分割任务设计，能确保图像和掩码的变换一致性。

3.2 损失函数与评估指标的选择

语义分割常用的损失函数是交叉熵损失（CrossEntropyLoss），特别是当你的掩码是单通道且每个像素值是类别索引时。PyTorch的nn.CrossEntropyLoss会自动处理ignore_index，这对于忽略某些无效标注区域（如标注图的边缘）非常有用。

criterion = nn.CrossEntropyLoss(ignore_index=255) # 假设255为忽略的标签

对于类别不平衡的数据集（例如，背景像素远多于前景），可以考虑使用Dice Loss或Focal Loss。Dice Loss直接优化分割区域的重叠度，对不平衡数据更鲁棒。我们可以结合两者：

class DiceLoss(nn.Module):
    def __init__(self, smooth=1e-6):
        super(DiceLoss, self).__init__()
        self.smooth = smooth

    def forward(self, pred, target):
        # pred: [B, C, H, W] 经过softmax或log_softmax
        # target: [B, H, W] 类别索引
        num_classes = pred.shape[1]
        # 将target转为one-hot格式 [B, C, H, W]
        target_one_hot = F.one_hot(target, num_classes).permute(0, 3, 1, 2).float()

        pred_softmax = F.softmax(pred, dim=1)

        intersection = (pred_softmax * target_one_hot).sum(dim=(2,3))
        union = pred_softmax.sum(dim=(2,3)) + target_one_hot.sum(dim=(2,3))

        dice = (2. * intersection + self.smooth) / (union + self.smooth)
        dice_loss = 1 - dice.mean()
        return dice_loss

# 组合损失
criterion_ce = nn.CrossEntropyLoss(ignore_index=255)
criterion_dice = DiceLoss()

def combined_loss(pred, target, alpha=0.5):
    ce_loss = criterion_ce(pred, target)
    dice_loss = criterion_dice(pred, target)
    return alpha * ce_loss + (1 - alpha) * dice_loss

评估指标方面，除了简单的像素准确率（Pixel Accuracy），更常用的是平均交并比（Mean Intersection over Union, mIoU），它能更好地衡量每个类别的分割质量。

def compute_iou(pred, target, n_classes, ignore_index=255):
    # pred: [B, H, W] 预测的类别索引
    # target: [B, H, W] 真实类别索引
    ious = []
    pred = pred.view(-1)
    target = target.view(-1)

    # 忽略特定索引
    valid = (target != ignore_index)
    pred = pred[valid]
    target = target[valid]

    for c in range(n_classes):
        pred_c = (pred == c)
        target_c = (target == c)
        intersection = (pred_c & target_c).sum().float()
        union = (pred_c | target_c).sum().float()
        if union == 0:
            ious.append(float('nan')) # 避免除零，此类别的IoU无定义
        else:
            ious.append((intersection / union).item())
    return np.nanmean(ious) # 计算均值时忽略NaN

3.3 训练循环与关键技巧

一个完整的训练循环需要整合数据加载、前向传播、损失计算、反向传播和优化器更新。这里给出一个训练epoch的核心逻辑框架：

def train_one_epoch(model, dataloader, criterion, optimizer, device, epoch, scheduler=None):
    model.train()
    total_loss = 0.0
    for batch_idx, (images, masks) in enumerate(dataloader):
        images, masks = images.to(device), masks.to(device)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, masks)

        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

        # 每N个batch打印一次日志
        if batch_idx % 50 == 0:
            print(f'Epoch [{epoch}], Batch [{batch_idx}/{len(dataloader)}], Loss: {loss.item():.4f}')

    avg_loss = total_loss / len(dataloader)
    if scheduler is not None:
        scheduler.step() # 按epoch调整学习率
    return avg_loss

在训练深度分割模型时，有几个技巧能显著提升效果和稳定性：

1. 学习率策略：使用CosineAnnealingLR或ReduceLROnPlateau调度器。初期可以用较小的学习率（如1e-4）预热（Warmup）几个epoch。
2. 优化器选择：AdamW（Adam with decoupled weight decay）通常是比原始Adam更好的选择，它能更有效地进行权重衰减。
3. 梯度裁剪：对于非常深的网络或大batch_size，梯度爆炸有时会发生。在loss.backward()之后，optimizer.step()之前加入torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)。
4. 混合精度训练：使用torch.cuda.amp进行自动混合精度训练，可以大幅减少显存占用并可能加快训练速度，尤其适用于显存紧张的情况。
5. 模型检查点与早停：定期保存验证集上性能最好的模型权重。如果连续多个epoch验证损失不再下降，则提前停止训练，防止过拟合。

4. 调试、可视化与模型部署

模型训练过程中，调试和可视化是理解模型行为、定位问题的关键。训练完成后，将模型投入实际应用则是最终目标。

4.1 常见问题调试指南

在实现和训练Deeplabv3+时，你可能会遇到以下典型问题：

• 问题：输出尺寸与输入尺寸不匹配。

  • 原因：上采样倍数计算错误。Backbone的下采样总倍数（通常是16或32）与解码器中上采样的倍数不匹配。
  • 解决：在模型forward函数的开始和每个上采样操作后，打印特征图的尺寸。确保最终上采样到(H, W)。

• 问题：损失为NaN或变得异常大。

  • 原因：学习率过高；数据中有异常值（如掩码标签超出类别范围）；网络层中出现了数值不稳定（如除零）。
  • 解决：
    1. 降低学习率（例如从1e-3降到1e-4）。
    2. 检查数据加载器，确保掩码标签值在[0, num_classes-1]或等于ignore_index。
    3. 在损失函数中加入平滑项（如Dice Loss中的smooth）。
    4. 使用梯度裁剪。

• 问题：训练时GPU内存溢出（OOM）。

  • 原因：batch_size太大；输入图像尺寸太大；模型参数量过大。
  • 解决：
    1. 减小batch_size。
    2. 在数据加载时对图像进行缩放（如缩放到512x512）。
    3. 考虑使用更轻量的Backbone（如ResNet18、MobileNetV2）。
    4. 启用梯度检查点（torch.utils.checkpoint）以时间换空间。
    5. 采用混合精度训练。

• 问题：模型在验证集上性能很差（过拟合）。

  • 原因：训练数据不足；模型复杂度太高；数据增强不够；训练时间太长。
  • 解决：
    1. 增强数据增强的强度（更多的随机裁剪、旋转、颜色扰动）。
    2. 在模型中增加Dropout层（如ASPP后、解码器卷积后）。
    3. 使用更激进的权重衰减。
    4. 采用早停策略。

4.2 训练过程可视化

使用TensorBoard或WandB等工具可以直观地监控训练过程。需要记录的信息包括：

• 训练损失和验证损失的变化曲线。
• 验证集mIoU和各类别IoU的变化。
• 学习率的变化。
• 训练和验证集的预测结果样例（图像、真实掩码、预测掩码）。

下面是一个使用TensorBoard记录图像样例的片段：

from torch.utils.tensorboard import SummaryWriter
import numpy as np

writer = SummaryWriter('runs/experiment_1')

def log_predictions(epoch, model, dataloader, device, num_images=3):
    model.eval()
    with torch.no_grad():
        images_logged = 0
        for images, masks in dataloader:
            if images_logged >= num_images:
                break
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)
            preds = torch.argmax(outputs, dim=1) # [B, H, W]

            for i in range(min(images.size(0), num_images - images_logged)):
                # 将Tensor转为numpy用于可视化
                img_np = images[i].cpu().numpy().transpose(1,2,0) # [H, W, C]
                mask_np = masks[i].cpu().numpy() # [H, W]
                pred_np = preds[i].cpu().numpy() # [H, W]

                # 可以创建一个叠加了预测边缘的彩色图像
                # ... 这里省略具体的可视化代码 ...
                vis_img = visualize_segmentation(img_np, mask_np, pred_np)

                # 添加到TensorBoard
                writer.add_image(f'Val_Sample_{images_logged}/Epoch_{epoch}', vis_img, epoch, dataformats='HWC')
                images_logged += 1
    model.train()

# 在验证循环后调用
# log_predictions(epoch, model, val_loader, device)

4.3 模型导出与部署

训练完成后，你可能需要将模型部署到生产环境。PyTorch提供了torch.jit.trace或torch.jit.script来生成TorchScript模型，或者使用ONNX格式进行跨平台部署。

导出为ONNX格式：

import torch.onnx

# 创建一个示例输入
dummy_input = torch.randn(1, 3, 512, 512).to(device)
model.eval()

# 导出模型
torch.onnx.export(
    model,
    dummy_input,
    "deeplabv3plus_resnet50.onnx",
    export_params=True,
    opset_version=12, # 选择一个合适的opset版本
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} # 支持动态batch
)
print("模型已导出为 ONNX 格式。")

在部署时，需要注意预处理（归一化、缩放）和后处理（argmax、上采样）需要与训练时保持一致。对于边缘设备，可以考虑使用模型量化和剪枝来减小模型体积和加速推理。PyTorch提供了torch.quantization模块支持动态和静态量化。

最后，一个完整的项目应该包含清晰的README.md，说明如何安装依赖、准备数据、训练模型以及进行推理。将配置参数（如学习率、批量大小、数据路径）外置到配置文件（如YAML）中，能让你的代码更易于维护和复现。记住，代码的清晰度和可复现性与模型性能同等重要。