PyTorch模型训练时未调用GPU、仅低负载使用CPU问题求助

1. 代码核心缺陷：gpu_ids参数未生效 #

你提供的代码虽然定义了--gpu_ids启动参数，但是没有对应的CUDA设备绑定逻辑，PyTorch没有被正确引导使用GPU资源，是GPU不生效的核心原因。

2. 第一步：先验证CUDA环境可用性 #

先在终端执行命令确认环境正常：

python -c "import torch; print(torch.cuda.is_available())"

如果返回False，说明你安装的是CPU版本PyTorch，或者CUDA驱动、环境配置错误，需要重新安装和你GPU驱动版本匹配的CUDA版PyTorch。
如果返回True，继续下一步修改代码。

3. 第二步：补全CUDA设备绑定逻辑 #

在main函数开头，拿到opt之后添加如下代码：

def main():
    opt = get_opt()
    # 新增设备绑定逻辑
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_ids
    if not torch.cuda.is_available():
        raise RuntimeError("未检测到可用GPU，请检查环境配置")
    # 原有代码继续
    print(opt)
    print("Start to train stage: %s, named: %s!" % (opt.stage, opt.name))
    ...

4. 第三步：解决CPU使用率低的问题 #

你当前启动参数设置的--workers 1，数据加载进程只有1个，很容易出现数据加载速度跟不上模型训练速度的瓶颈，导致CPU、GPU使用率都偏低。你可以根据你的CPU核心数调整这个参数，比如改成4~8，修改后的启动命令参考：

python train.py --name GMM --stage GMM --workers 4 --save_count 5000 --shuffle --gpu_ids 0

5. 补充注意事项 #

如果后续训练TOM阶段，需要修改train_tom函数里的变量，当前代码里im_pose、im_h、shape三个输入没有调用.cuda()转移到GPU，会导致数据和模型不在同一个设备报错，需补充对应的转移逻辑。

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