### 使用 PyTorch 进行语义分割的批量预测 在使用 PyTorch 实现语义分割的过程中，模型训练完成后通常会进入测试阶段，在此阶段可以利用批量输入来提高推理效率。以下是实现批量预测的核心要点： #### 数据预处理 为了进行批量预测，需要确保输入图像经过标准化和尺寸调整等操作后能够形成一个张量批次。这可以通过 `torch.utils.data.DataLoader` 来完成。该工具支持自动批量化输入数据。 ```python from torch.utils.data import DataLoader # 假设已经定义了一个自定义的数据集类 SegmentationDataset data_loader = DataLoader( dataset=SegmentationDataset(...), # 替换为实际参数 batch_size=8, shuffle=False, num_workers=4 ) ``` 以上代码片段展示了如何创建一个加载器用于批量读取数据[^1]。 #### 模型推断 当数据被分批加载后，可以直接将其送入已训练好的模型中执行前向传播运算。注意设置模型为评估模式 (`model.eval()`) 并禁用梯度跟踪以节省内存资源。 ```python import torch device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(device) with torch.no_grad(): model.eval() for images, _ in data_loader: images = images.to(device) outputs = model(images) # 输出维度应为 [batch_size, num_classes, height, width] _, predicted_masks = torch.max(outputs, dim=1) # 获取每像素的最大概率类别索引 # 将预测掩码转回 CPU 或保存到磁盘 for i, mask in enumerate(predicted_masks.cpu().numpy()): save_mask(mask, f"mask_{i}.png") # 自定义函数保存图片 ``` 这里的关键在于通过调用 `torch.max()` 函数沿通道轴获取最高置信度对应的类别标签作为最终预测结果[^3]。 #### 后处理与可视化 得到预测后的掩膜可能还需要一些额外的操作比如重采样回到原始分辨率或者应用颜色映射以便于观察效果。 ```python def visualize_predictions(image_tensor, pred_mask): image_np = np.transpose(image_tensor.squeeze(0).cpu().numpy(), (1, 2, 0)) plt.figure(figsize=(10,5)) plt.subplot(1,2,1) plt.imshow(image_np.astype(np.uint8)) # 显示原图 plt.subplot(1,2,2) color_map = get_color_map(pred_mask.shape[-2:]) # 定义自己的配色方案 colored_pred = apply_colormap_to_image(color_map, pred_mask.numpy()) plt.imshow(colored_pred) # 展示着色后的预测结果 plt.show() for imgs, preds in zip(batch_images[:], predicted_masks[:]): visualize_predictions(imgs, preds) ``` 上述脚本提供了一种简单的方法展示单幅或多幅样本及其对应预测值之间的对比关系[^2]。