如何在SageMaker实时推理中利用全部GPU？代码修改咨询

一、核心思路 #

通过PyTorch多进程（torch.multiprocessing）为每个GPU分配独立进程，每个进程单独加载一套模型到对应GPU，实现并行推理，缩短单批次或多请求的处理时间。

二、具体代码修改 #

1. 模型加载逻辑改造 #

移除全局DEVICE变量，改为在进程内绑定指定GPU，每个进程独立加载模型，避免跨进程CUDA上下文冲突。

def load_model(rank, model_dir, queue):
    # 绑定当前进程到目标GPU
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")
    
    logger.info(f"加载第一个模型到GPU {rank}...")
    model = Model().to(device)
    with open(os.path.join(model_dir, "checkpoint.pth"), "rb") as f:
        model.load_state_dict(torch.load(f, map_location=device)['state_dict'])
    model = model.eval()
    
    logger.info(f"加载第二个模型到GPU {rank}...")
    model_2 = Model_2()
    model_2.to(device)
    checkpoint = torch.load('checkpoint_2.pth', map_location=device)
    # 修正原代码笔误：此处应为load_state_dict而非直接调用模型
    model_2.load_state_dict(remove_prefix_state_dict(checkpoint['state_dict']), strict=True)
    model_2 = model_2.eval()
    
    logger.info(f"GPU {rank}模型加载完成")
    # 将加载好的模型存入队列，供主进程调度
    queue.put((rank, {'first_model': model, 'second_model': model_2}))

2. 多进程模型初始化 #

在服务启动阶段，启动与GPU数量匹配的进程，完成所有模型的加载：

import torch.multiprocessing as mp

def init_multigpu_models(model_dir):
    num_gpus = torch.cuda.device_count()
    if num_gpus == 0:
        raise RuntimeError("未检测到可用GPU")
    
    # 设置进程启动方式为spawn（PyTorch CUDA多进程强制要求）
    mp.set_start_method('spawn', force=True)
    
    queue = mp.Queue()
    processes = []
    
    for rank in range(num_gpus):
        p = mp.Process(target=load_model, args=(rank, model_dir, queue))
        p.start()
        processes.append(p)
    
    # 收集所有进程加载的模型
    model_dict = {}
    for _ in range(num_gpus):
        rank, models = queue.get()
        model_dict[rank] = models
    
    # 等待所有进程初始化完成
    for p in processes:
        p.join()
    
    return model_dict, num_gpus, processes

3. 输入处理优化 #

支持批量请求输入，方便拆分到多个GPU并行处理，同时兼容单请求场景：

def input_fn(request_body, request_content_type):
    assert request_content_type=='application/json'
    req_data = json.loads(request_body)
    
    # 兼容单URL/批量URL两种输入格式
    urls = req_data.get('urls', [req_data['url']])
    save_names = req_data.get('save_names', [req_data['save_name']])
    
    input_tensors = []
    metadata_list = []
    
    for url, save_name in zip(urls, save_names):
        logger.info(f'处理图片URL: {url}')
        img = Image.open(requests.get(url, stream=True).raw).convert('RGB')
        w, h = img.size
        input_tensor = preprocess(img)
        
        input_tensors.append(input_tensor)
        metadata_list.append({'w':w,'h':h,'image':img, 'save_name':save_name})
    
    # 合并为批量张量，分配时再转至对应GPU
    input_batch = torch.stack(input_tensors).to('cpu')
    return {'tensor': input_batch, 'metadata': metadata_list}

4. 预测逻辑与GPU调度 #

实现轮询调度，将请求拆分到不同GPU并行处理：

from itertools import cycle

# 全局GPU调度器，轮询分配请求
gpu_scheduler = None

def predict_fn(input_object, model_dict):
    global gpu_scheduler
    if gpu_scheduler is None:
        gpu_scheduler = cycle(model_dict.keys())
    
    input_batch = input_object['tensor']
    metadata_list = input_object['metadata']
    batch_size = input_batch.shape[0]
    num_gpus = len(model_dict)
    
    # 将批量输入拆分为与GPU数量匹配的子批量
    split_batches = torch.chunk(input_batch, num_gpus)
    results = []
    
    for batch_idx, sub_batch in enumerate(split_batches):
        rank = next(gpu_scheduler)
        device = torch.device(f"cuda:{rank}")
        model_1 = model_dict[rank]['first_model']
        model_2 = model_dict[rank]['second_model']
        
        # 将子批量移至目标GPU
        sub_batch = sub_batch.to(device)
        
        # 模型推理（禁用梯度计算加速）
        with torch.no_grad():
            d0, d1, d2, d3, d4, d5, d6 = model_1(sub_batch)
        
        # 处理推理结果，移回CPU
        masks = torch.argmax(d0, axis=1).cpu().numpy()
        for mask_idx, mask in enumerate(masks):
            # 匹配对应元数据
            meta = metadata_list[batch_idx*len(masks)+mask_idx]
            
            mask = np.where(mask==2, 255, mask)
            mask = np.where(mask==1, 128, mask)
            final_image = Image.fromarray(mask).resize((meta['w'], meta['h'])).convert('L')
            
            img = np.array(meta['image'])[:,:,::-1]
            final_image = np.array(final_image)
            image_dict = to_dict(img, final_image)
            final_image = model_2_process(model_2, image_dict)
            
            results.append({
                "final_output": final_image, 
                'image': meta['image'], 
                'save_name': meta['save_name']
            })
    
    return results

三、多进程实现关键注意事项 #

• 进程启动方式：必须使用spawn或forkserver，禁止使用fork，否则会导致CUDA上下文混乱
• 模型独立性：每个进程必须独立加载模型，禁止跨进程共享模型参数
• 内存管理：仅在必要时调用torch.cuda.empty_cache()，过度调用会降低推理性能
• 调度策略：轮询适合批量请求场景，单请求场景可改用空闲GPU队列调度
• 异常防护：为每个进程添加异常捕获逻辑，避免单个进程崩溃导致整个服务宕机
• 权重加载修正：原代码中model_2的权重加载存在笔误，必须改为load_state_dict才能正确加载权重

内容的提问来源于stack exchange，提问作者Diego Rodea