PyTorch UNet分割模型多GPU推理仅单卡运行问题求助

核心问题：单卡闲置的原因 #

你当前推理时batch size=1，而DataParallel的工作逻辑是将一个batch的样本拆分到多张GPU并行计算。当batch仅含1个样本时，无法拆分分配，只会占用第一张GPU，另一卡闲置，最终导致单卡内存过载触发OOM。

具体修复步骤 #

1. 改为批量推理，适配多卡并行 #

将多张图像打包成batch输入，让DataParallel能把样本分配到两张GPU上。修改代码如下：

import torch
import cv2
import numpy as np
from model import build_unet
from torch.nn.parallel import DataParallel
import os
from tqdm import tqdm
from torch.cuda.amp import autocast  # 用于混合精度

torch.cuda.empty_cache()
checkpoint_path = "Weights/best_model.pth"

def mask_parse(mask):
    mask = np.expand_dims(mask, axis=-1)
    mask = np.concatenate([mask, mask, mask], axis=-1)  
    return mask

if __name__ == "__main__":
    model = build_unet()
    device = torch.device("cuda")  
    model = model.to(device)
    model = DataParallel(model)
    # 处理预训练权重的键名适配
    checkpoint = torch.load(checkpoint_path, map_location="cuda")
    if not list(checkpoint.keys())[0].startswith('module.'):
        checkpoint = {f'module.{k}': v for k, v in checkpoint.items()}
    model.load_state_dict(checkpoint)
    model.eval()  # 启用推理模式，禁用训练冗余操作
    
    dirListCSV = os.listdir('TestImages')
    batch_size = 2  # 刚好适配2张GPU，每张卡处理1个样本

    # 批量处理图像
    for i in tqdm(range(0, len(dirListCSV), batch_size)):
        batch_files = dirListCSV[i:i+batch_size]
        batch_imgs = []
        batch_meta = []  # 保存原图尺寸、文件名

        for allCSV in batch_files:
            imgName = 'TestImages/'+ allCSV
            cv_img = cv2.imread(imgName)
            mHeight, mWidth = cv_img.shape[0], cv_img.shape[1]
            cv_img = cv2.resize(cv_img, (1024, 1792))
            batch_imgs.append(cv_img)
            batch_meta.append((mHeight, mWidth, allCSV))
        
        # 转换为模型输入格式
        x = np.array(batch_imgs)
        x = np.transpose(x, (0, 3, 1, 2))  # 转为(batch, channel, h, w)
        x = x / 255.0
        x = x.astype(np.float32)
        x = torch.from_numpy(x).to(device)
        
        # 推理阶段
        with torch.no_grad(), autocast():  # 混合精度减少内存占用
            pred_y = model(x)
            pred_y = torch.sigmoid(pred_y)
            
            # 逐个处理batch结果
            for idx in range(pred_y.shape[0]):
                mask = pred_y[idx].cpu().numpy()
                mask = np.squeeze(mask, axis=0)
                mask = mask > 0.1
                mask = np.array(mask, dtype=np.uint8)
                mask = mask_parse(mask) * 255
                
                # 恢复原图尺寸并保存
                mH, mW, imname = batch_meta[idx]
                original_img = cv2.imread(f'TestImages/{imname}')
                original_img = cv2.resize(original_img, (mW, mH))
                mask = cv2.resize(mask, (mW, mH))
                finalOut = cv2.addWeighted(original_img, 0.6, mask, 0.4, 0)
                cv2.imwrite(f'Output/{imname}', finalOut)
        
        # 打印内存使用并清理缓存
        print('Gpu0:',torch.cuda.max_memory_allocated(device=0))
        print('Gpu1:',torch.cuda.max_memory_allocated(device=1))
        torch.cuda.empty_cache()

2. 额外内存优化细节 #

• 启用autocast混合精度推理：能在几乎不损失精度的前提下，将张量内存占用减半。
• 确保调用model.eval()：禁用训练时的BatchNorm统计更新、Dropout等操作，减少内存开销和计算量。
• 批量后强制清理缓存：避免张量堆积占用内存。

效果验证 #

修改后，两张GPU会各自处理batch中的一个样本，内存压力被分摊，OOM问题会得到缓解，同时批量推理也能提升整体处理速度。

内容的提问来源于stack exchange，提问作者vyshak balakrishnan