如何解决Colab中PyTorch大尺寸图像分割的CUDA内存不足问题？

针对你在2448×2448分辨率图像分割预测时遇到的CUDA内存溢出问题，以下是实用的解决手段：

1. 优化PyTorch内存分配策略 #

在代码最开头添加环境变量配置，解决显存碎片问题：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

2. 降低单张图像的显存占用 #

临时缩放输入图像 #

将高分辨率图像缩小后再预测，最后把结果放大回原尺寸（适合对精度要求不是极端严格的场景）：

from torchvision.transforms import Resize

# 推理前先缩小输入
resize_input = Resize((1224, 1224))
x_tensor_small = resize_input(x_tensor)

with torch.no_grad():
    pred_mask_small = best_model(x_tensor_small)

# 将预测结果放大回原分辨率
resize_output = Resize((2448, 2448))
pred_mask = resize_output(pred_mask_small)

分块预测大图像 #

把2448×2448的图像切割成多个小patch（比如4个1224×1224块），分别预测后拼接成完整结果，避免一次性加载超大张量：

# 示例：将图像切为2x2的块
patch_size = 1224
pred_patches = []
for i in range(2):
    row_patches = []
    for j in range(2):
        # 提取当前patch
        patch = x_tensor[:, :, i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size]
        with torch.no_grad():
            patch_pred = best_model(patch)
        row_patches.append(patch_pred)
    # 拼接行内的patch
    row_pred = torch.cat(row_patches, dim=3)
    pred_patches.append(row_pred)
# 拼接所有行得到完整预测结果
pred_mask = torch.cat(pred_patches, dim=2)

3. 强制禁用梯度计算与及时清理显存 #

• 用torch.no_grad()包裹推理代码，彻底关闭梯度计算（训练时的梯度会占用大量显存）：

  with torch.no_grad():
      pred_mask = best_model(x_tensor)
  

• 循环内每处理完一张图像，手动清理无用张量并释放显存缓存：

  # 处理完当前图像后执行
  del x_tensor, pred_mask, gt_mask
  torch.cuda.empty_cache()
  

4. 确保模型处于评估模式 #

推理前必须将模型切换到评估模式，避免BatchNorm等层保留训练状态占用额外内存：

best_model.eval()  # 放在循环前执行一次即可

内容的提问来源于stack exchange，提问作者Om Soni