Win10中JIT调用自定义CUDA算子踩坑

小潮哥

已于 2022-08-18 12:00:20 修改

阅读量3.6k

点赞数 2

分类专栏：环境配置文章标签： pytorch

于 2022-08-17 22:01:18 首次发布

本文链接：https://blog.csdn.net/ytrewq121/article/details/126394736

版权

环境配置专栏收录该内容

3 篇文章 1 订阅

订阅专栏

这篇博客介绍了在Windows环境下使用PyTorch自定义C++扩展时遇到的常见问题，包括'ninja'命令错误、找不到指定模块以及'torch.utils.cpp_extension.load'卡住无响应的解决办法。通过设置环境变量、手动编译'ninja'和处理lock文件，成功编译并加载自定义模块。

摘要由CSDN通过智能技术生成

硬件、系统与软件

WIN10 + Visual Studio2019(NOTE: 我的VS2019装在D:\下)

Nvidia RTX 3070 Laptop NVIDIA-SMI 512.33 Driver Version: 512.33 CUDA Version: 11.3

Pycharm + Anaconda + python3.8 + pytorch-gpu

与JIT相关的环境变量

Path

C:\local\Ninja
D:\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.29.30133
D:\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.29.30133\bin\Hostx64\x64

INCLUDE

D:\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.29.30133\include
C:\Program Files (x86)\Windows Kits\10\Include\10.0.19041.0\shared
C:\Program Files (x86)\Windows Kits\10\Include\10.0.19041.0\ucrt
C:\Program Files (x86)\Windows Kits\10\Include\10.0.19041.0\um
C:\Program Files (x86)\Windows Kits\10\Include\10.0.19041.0\winrt

LIB

D:\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.29.30133\lib\x64
C:\Program Files (x86)\Windows Kits\10\Lib\10.0.19041.0\um\x64
C:\Program Files (x86)\Windows Kits\10\Lib\10.0.19041.0\ucrt\x64
C:\Program Files (x86)\Microsoft SDKs\Windows\v7.1A\Lib

遇到的问题和bug

1.Command '['ninja', '-v']' returned non-zero exit status 1

首先要安装ninja并设置环境变量(前文PATH中有)。其次千万不要把'-v'改成'--v'或者'--version'作为彻底的解决方案，否则JIT将无法自动调用ninja进行编译，属于治标不治本的行为。可以暂时改成'--verision'，然后在将来能够生成build.ninja文件的时候再改回去。但此处虽然报错并不是错误。

2.找不到指定模块

运行python包含load()的主函数后，发现下面错误

module = importlib.util.module_from_spec(spec)

File "<frozen importlib._bootstrap>", line 556, in module_from_spec

File "<frozen importlib._bootstrap_external>", line 1166, in create_module

File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed

这是由于那个模块并没有被编译生成有效的pyd文件。虽然编译没有成功，但是类似于CMakeLists.txt的build.ninja成功生成了。此时再继续看python的报错信息很难看出修改方案。根据pycharm输出的信息，找到build.ninja文件。

D:\Anaconda3\envs\gpupy\python.exe E:/wjc_code/python/ctorch_sample_3/time.py
Using C:\Users\85063\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu113 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file C:\Users\85063\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu113\add2\build.ninja...

开启一个命令行，进入build.ninja所在文件夹，手动编译这个build.ninja。

C:\Users\85063\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu113\add2>ninja

过程中会出现一些找不到文件"xxx.lib"之类的错误，通过前文中设置好环境变量的方式解决他们。编译成功后应该会生成一个pyd文件，这就是python可以直接理解的库。

3.torch.utils.cpp_extension.load卡住无响应

D:\Anaconda3\envs\gpupy\python.exe E:/wjc_code/python/ctorch_sample_3/time.py
Using C:\Users\85063\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu113 as PyTorch extensions root...

###然后再没反应了

查看任务管理器，发现新进程完全不占用CPU和内存。这是由于曾经同时开启了两个程序调用这个算子，导致生成了一个lock文件，阻塞了其他进程使用自定义的算子。手动将其删除即可。或者参考

torch.utils.cpp_extension.load卡住无响应_zParquet的博客-CSDN博客_torch.utils.cpp_extension.load