PyTorch中net.cuda()或net.to(device)执行耗时过长原因

正常情况下模型迁移到GPU的耗时仅和模型总大小正相关，小型模型耗时应该在毫秒级，卡30分钟一定是环境或模型定义异常导致，按以下步骤排查：

第一步：先定位是CUDA环境问题还是模型本身问题 #

在执行模型迁移代码前，先运行最小CUDA测试代码：

import torch
print(torch.cuda.is_available())
test_tensor = torch.tensor([1.0]).cuda()
print("CUDA上下文初始化完成")

• 如果这几行代码就需要很久才能跑完，问题出在CUDA运行环境，和模型无关：

  • 检查笔记本电源状态：必须插外接电源，系统电源模式开到「最佳性能」，不插电时笔记本RTX3050会处于低功耗休眠状态，第一次唤醒+初始化CUDA上下文可能出现超长阻塞。
  • 切换显卡输出模式：在笔记本自带的性能控制中心把显卡输出从「混合模式」改成「独显直连」，重启电脑后再试。混合输出模式下核显负责画面输出、独显负责计算，WDDM调度层经常出现兼容问题导致CUDA调用阻塞。
  • 关闭Windows硬件加速GPU计划：打开「设置-系统-显示-图形-默认图形设置」，关闭「硬件加速GPU计划」选项后重启，该功能和30系笔记本GPU、PyTorch的兼容问题非常普遍，是CUDA初始化卡顿的高发原因。
  • 重装匹配版本的PyTorch：卸载当前环境下的PyTorch，安装对应CUDA11.8的稳定正式版PyTorch，不要装nightly预览版，也不要装和当前显卡驱动版本不匹配的高版本CUDA对应PyTorch，版本不兼容时驱动会在第一次CUDA调用时做大量JIT重编译、兼容性校验，耗时可能长达数十分钟。
  • 清理显存占用：打开任务管理器查看RTX3050的专用显存占用，关闭所有占显存的无关程序（比如开了硬件加速的浏览器、游戏、视频剪辑软件），至少留出2GB以上空闲显存再跑代码。

• 如果上述测试代码可以秒级跑完，问题出在模型定义本身：

  • 先统计模型实际总大小，运行以下代码计算参数+缓冲区的总内存占用：

    total_mem = 0
    for param in net.parameters():
        total_mem += param.numel() * param.element_size()
    for buffer in net.buffers():
        total_mem += buffer.numel() * buffer.element_size()
    print(f"模型实际占用大小: {total_mem / 1024 / 1024:.2f} MB")
    

    正常小型模型的大小应该在几MB到几百MB区间，如果输出结果超过2GB，说明误将超大张量（比如完整数据集、超大查找表）注册成了模型的参数或缓冲区，.cuda()执行时需要把这些大张量全部拷贝进仅4GB容量的显存，显存不足时WDDM会在内存和显存之间反复做页交换，就会出现数十分钟的卡顿。把这类不属于模型权重的大张量从模型定义中移除，用到时再单独拷贝到GPU即可。
  • 检查模型__init__方法的自定义逻辑：如果写了自定义CUDA算子初始化、设备相关的阻塞等待逻辑，这类逻辑在模型迁移设备时会被触发，把相关初始化逻辑挪到模型迁移到GPU之后再执行即可。

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