双GPU环境下Flask中torch.load指定cuda:1报错的原因与解决方法

我之前碰到过好几个类似的Flask+PyTorch多GPU问题，你的情况核心原因是Flask的多进程/多线程启动机制破坏了PyTorch的GPU设备上下文继承，具体拆解下：

为什么会出现这个矛盾？ #

• 你直接跑torch.cuda.device_count()是单进程环境，系统能正确识别所有2块GPU；但Flask默认的开发服务器（比如Werkzeug）会自动以多进程模式启动，当它fork子进程时，子进程无法完整继承父进程的CUDA设备上下文，导致子进程只能识别到1块GPU（极端情况甚至识别不到）。
• 训练阶段你是单进程运行模型，没有进程fork的问题，所以torch.load(map_location='cuda:1')能正常找到设备；但Flask的子进程里torch.cuda.device_count()变成了1，这时候你要加载到cuda:1就会报错“设备2不存在”（这里的设备编号是子进程里的视角，相当于子进程只认cuda:0，你指定的cuda:1在它眼里就是不存在的设备2）。
• 另外还有一种可能：Flask启动时意外修改了CUDA_VISIBLE_DEVICES环境变量，导致子进程只能看到部分GPU。

具体解决办法 #

1. 强制Flask用单进程模式启动 #

这是最直接的解决方式，避免进程fork带来的上下文问题：

• 命令行启动时加参数：

  flask run --without-threads
  

• 或者在代码里显式设置：

  if __name__ == '__main__':
      app.run(threaded=False, processes=1)
  

  单进程模式下，Flask进程能完整继承GPU设备列表，torch.cuda.device_count()就会正常返回2了。

2. 手动锁定CUDA可见设备 #

在Flask应用的最开头（一定要在导入torch之前）手动设置环境变量，确保所有进程都能看到两块GPU：

import os
# 强制让进程可见两块GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
import torch

这样不管Flask怎么启动，子进程都会继承这个环境变量，不会出现GPU丢失的情况。

3. 用动态映射函数替代固定设备编号 #

不要直接写死map_location='cuda:1'，而是用函数动态适配当前可用的GPU：

def dynamic_map_location(storage, loc):
    # 如果有至少2块GPU就用cuda:1，否则用cuda:0
    target_device = 1 if torch.cuda.device_count() >= 2 else 0
    return storage.cuda(target_device)

model = torch.load(model_path, map_location=dynamic_map_location)

这种方式更鲁棒，就算GPU环境变化也不会报错。

4. 生产环境服务器的适配 #

如果用Gunicorn等生产服务器，注意不要开过多worker，或者给每个worker绑定单独的GPU：

• 单worker启动（避免多进程冲突）：

  CUDA_VISIBLE_DEVICES=0,1 gunicorn --workers=1 app:app
  

• 如果需要多worker，给每个worker分配不同的GPU：

  # 第一个worker用cuda:0
  CUDA_VISIBLE_DEVICES=0 gunicorn --workers=1 --bind 0.0.0.0:5000 app:app &
  # 第二个worker用cuda:1
  CUDA_VISIBLE_DEVICES=1 gunicorn --workers=1 --bind 0.0.0.0:5001 app:app &
  

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