pytorch 语义分割 训练自己的数据集

### 构建自定义数据集并进行语义分割训练 为了使用 PyTorch 进行自定义数据集上的语义分割训练，需遵循一系列流程来准备数据、构建模型和设置训练环境。 #### 数据集准备 创建自定义的数据加载器是至关重要的一步。这通常涉及到继承 `torch.utils.data.Dataset` 类，并重写其方法以适应特定需求。对于视网膜血管分割这样的任务，可以参照如下结构： ```python from torchvision import transforms as T import os.path as osp from PIL import Image import numpy as np class MyDataset(torch.utils.data.Dataset): CLASSES = ['background', 'vessel'] PALETTE = [[0, 0, 0], [255, 255, 255]] def __init__(self, data_root, img_dir, mask_dir, transform=None): self.transform = transform self.imgs = list(sorted(os.listdir(osp.join(data_root, img_dir)))) self.masks = list(sorted(os.listdir(osp.join(data_root, mask_dir)))) def __getitem__(self, idx): # 加载图片与对应的mask img_path = osp.join(self.root, "images", self.imgs[idx]) mask_path = osp.join(self.root, "masks", self.masks[idx]) image = Image.open(img_path).convert('RGB') target = Image.open(mask_path) if self.transforms is not None: transformed = self.transforms(image=np.array(image), mask=np.array(target)) image = transformed['image'] target = transformed['mask'] return image, target def __len__(self): return len(self.imgs) ``` 此代码片段展示了如何通过继承 `Dataset` 来建立一个新的类 `MyDataset`[^3]，该类能够处理来自指定文件夹的图像及其相应的掩码（即标签）。注意这里假设输入图像是三通道彩色图像而掩码为灰度级表示的目标区域。 #### 创建 U-Net 模型架构 接下来就是搭建适合于所选任务的神经网络框架，在这个例子中选择了经典的 U-Net 结构作为基础模型来进行二分类——区分背景和其他感兴趣的对象（比如血管）: ```python import torch.nn.functional as F from collections import OrderedDict import torch.nn as nn def double_conv(in_channels, out_channels): return nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3), nn.ReLU(inplace=True)) class UNet(nn.Module): def __init__(num_classes=1): # 对于二元分类，默认类别数设为1 super().__init__() self.dconv_down1 = double_conv(3, 64) ... # 定义完整的编码解码过程... self.conv_last = nn.Conv2d(64, num_classes, 1) def forward(x): conv1 = self.dconv_down1(x) ... # 编码解码的具体操作... output = self.conv_last(output) return output ``` 上述代码仅提供了一个简化版的 U-Net 实现方式；实际应用时可能还需要考虑更多细节如跳跃连接的设计等[^1]。 #### 设置训练参数及优化策略 最后要配置好训练过程中需要用到的各种超参数以及选择合适的损失函数和优化算法： ```python device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = UNet().to(device=device) criterion = BCEWithLogitsLoss() optimizer = Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): model.train() running_loss = .0 for images, masks in dataloader: optimizer.zero_grad() outputs = model(images.to(device)) loss = criterion(outputs.squeeze(), masks.float().to(device)) loss.backward() optimizer.step() running_loss += loss.item() ``` 这段伪代码说明了在一个典型的训练循环里应该做些什么工作：初始化设备、实例化模型对象并将之移动到 GPU 上运行（如果有可用的话）、设定交叉熵或其他适当类型的代价函数还有梯度下降法求解最优解等等[^2]。