如何正确在GPU上运行DataWig？传GPU上下文后仍运行在CPU怎么办

• 调用DataWig的Imputer.fit方法时，已按照MXNet规范传入指定GPU ID的上下文参数，指定运行设备为GPU，但训练流程实际仍在CPU上执行
• 当前环境的CUDA 11.0版本不在MXNet官方原生支持的版本匹配列表内
• 对应的参考调用代码：

from mxnet import gpu
...
ctx = gpu(0)
...
imputer.fit(train_df=imputer_train_df, num_epochs=10, ctx=ctx)

• 完成MXNet与CUDA版本的匹配校准
  • 不要在CUDA 11.0环境下安装纯CPU版本的MXNet包，MXNet的GPU能力和安装的构建包后缀严格绑定，CUDA 11.0对应的官方可用MXNet构建版本为1.7.x系列带cu110标识的安装包，直接卸载原有CPU版MXNet，安装对应匹配的mxnet-cu110包即可完成基础版本对齐，否则即使传入GPU上下文参数，MXNet也会自动回退到CPU设备执行
• 提前校验MXNet的GPU可访问性
  • 安装完成后优先运行最简测试代码，确认MXNet可正常调用GPU：

    from mxnet import gpu, ndarray
    # 尝试在GPU上生成测试数组，无报错即代表GPU上下文可正常被MXNet识别
    test_tensor = ndarray.random.normal(shape=(2,2), ctx=gpu(0))
    print(test_tensor.context)
    

  • 若测试过程中抛出无GPU设备、上下文不匹配类报错，说明安装的MXNet构建包和当前CUDA版本不兼容，需要重新安装版本匹配的对应包
• 适配DataWig的参数透传逻辑
  • 部分旧版本的DataWig存在参数兼容问题，即使在fit方法中传入了自定义GPU上下文参数，也不会将参数透传到内部的MXNet训练循环，会强制使用默认CPU上下文，需确认当前安装的DataWig版本为0.1.10及以上，该版本才正式支持在fit方法中接收并透传自定义MXNet上下文
  • 若使用的DataWig版本无法升级，可通过全局上下文覆盖的方式实现GPU运行，无需依赖方法参数传值：

    import os
    os.environ["MXNET_GPU_WORKER_NTHREADS"] = "2"
    # 将MXNet全局默认上下文强制设置为指定GPU
    from mxnet import context, gpu
    context.Context._default = gpu(0)
    

• 实际运行状态校验
  • 训练流程启动后，可通过nvidia-smi命令查看GPU显存占用情况，若对应训练进程出现明确的GPU显存占用增长，即代表程序实际运行在GPU上，无需仅依赖代码内的参数打印判断运行设备

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