bug探讨与解决：StyleGAN2和StyleGAN3 ImportError: DLL load failed

问题描述

笔者的环境

• OS: Windows 11
• Python: 3.8
• PyTorch: 1.10.1 (built with cuda 11.3)
• torchvision: 0.11.2
• CUDA toolkit: 12.1
• ninja: 1.11.1.git.kitware.jobserver-1
• Visual Studion 2022
• cl.exe: 19.35.32215
• Nvidia Driver Version: 531.29

这里其实可以看见，cuda toolkit版本和torch使用的版本不一样，这里是个伏笔，划重点

出现的问题

Traceback (most recent call last):
  File ".\pvg_lhq-test.py", line 21, in <module>
    from models.render_model import RenderModel
  File "C:\Users\**\workspace\google-research\infinite_nature_zero\models\render_model.py", line 29, in <module>
    from models.networks import co_mod_gan
  File "C:\Users\**\workspace\google-research\infinite_nature_zero\models\networks\co_mod_gan.py", line 23, in <module>
    from models.networks.model import ConvLayer
  File "C:\Users\**\workspace\google-research\infinite_nature_zero\models\networks\model.py", line 11, in <module>
    from .op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d, conv2d_gradfix
  File "C:\Users\**\workspace\google-research\infinite_nature_zero\models\networks\op\__init__.py", line 1, in <module>
    from .fused_act import FusedLeakyReLU, fused_leaky_relu
  File "C:\Users\**\workspace\google-research\infinite_nature_zero\models\networks\op\fused_act.py", line 11, in <module>
    fused = getmodule(
  File "C:\Users\**\workspace\google-research\infinite_nature_zero\tools\module.py", line 23, in getmodule
    return load(name, sources, **kargs)
  File "C:\Users\**\.conda\envs\infinite_nature_zero\lib\site-packages\torch\utils\cpp_extension.py", line 1124, in load
    return _jit_compile(
  File "C:\Users\**\.conda\envs\infinite_nature_zero\lib\site-packages\torch\utils\cpp_extension.py", line 1362, in _jit_compile
    return _import_module_from_library(name, build_directory, is_python_module)
  File "C:\Users\**\.conda\envs\infinite_nature_zero\lib\site-packages\torch\utils\cpp_extension.py", line 1752, in _import_module_from_library
    module = importlib.util.module_from_spec(spec)
ImportError: DLL load failed while importing fused: 找不到指定的模块。

另外还有可能出现类似于：

C:\Users\**\.conda\envs\infinite_nature_zero\lib\site-packages\torch\utils\cpp_extension.py" line 1682 code from "command = ['ninja', '-v']" to " command = ['ninja', '--version']

问题分析

分析步骤和思路

1. 通过torch.utils.cpp_extension._get_build_directory(name, False)函数可以获得ninja的工作路径，以fused模块为例，torch.utils.cpp_extension._get_build_directory(‘fused', False)返回的路径下面有编译生成的二进制文件ninja.pyd、ninja文件build.ninja等。我们查看build.ninja：

ldflags = /DLL c10.lib c10_cuda.lib torch_cpu.lib torch_cuda_cu.lib -INCLUDE:?searchsorted_cuda@native@at@@YA?AVTensor@2@AEBV32@0_N1@Z torch_cuda_cpp.lib -INCLUDE:?warp_size@cuda@at@@YAHXZ torch.lib /LIBPATH:C:\Users\frank\.conda\envs\infinite_nature_zero\lib\site-packages\torch\lib torch_python.lib /LIBPATH:C:\Users\frank\.conda\envs\infinite_nature_zero\libs "/LIBPATH:C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1\lib/x64" cudart.lib

这里可以看见连接的动态链接库包括cudart等，这里记一下
2. 通过Visual Studio等工具对运行时加载的dll文件进行分析，发现PyTorch加载的cudart为自带的cudart64_110.dll，而通过dumpbin.exe /dependents C:\Users\**\AppData\Local\torch_extensions\torch_extensions\Cache\py38_cu113\fused\fused.pyd得知，fused.pyd依赖的是cuda 12.1的cudart64_12.dll (dumpbin.exe为VS的开发者命令行工具，使用的时候需要先运行Developer Command Prompt for VS 2022或者Developer PowerShell for VS 2022来配置环境)，由此可见找不到的DLL为pyd文件所依赖版本的cudart动态链接库。
3. 通过检查环境变量发现，cuda 12.1在环境变量中，在C++中以下代码可以正常运行，且fused.pyd加载正常，并未报错

#include <Windows.h>
#include <Python.h>
#include <iostream>

// 此处为笔者机器里面pyd的生成路径
auto pydpath = L"C:\\Users\\**\\AppData\\Local\\torch_extensions\\torch_extensions\\Cache\\py38_cu113\\fused\\fused.pyd";

typedef void* init();

int main()
{
	// 必须先初始化python才能调用python模块
	Py_Initialize();
	auto dll = LoadLibrary(pydpath);
	auto err = GetLastError();
	if (dll) 
	{
		auto address = GetProcAddress(dll, "PyInit_fused");
		// 调用初始化函数
		((void(*)())address)();
	}
	else
	{
		std::cout << err << std::endl;
	}
	Py_Finalize();
	return 0;
}

分析结果

综上所述，我们可以得出以下结论：

1. 该类型的bug产生的原因为jit使用的依赖和pytorch等同名依赖版本不同，比如笔者就是因为pytorch使用的是cuda 11.3的cudart64_110.dll，而fused.pyd使用的是cuda 12.1的cudart64_12.dll。
2. python3.8 (3.7测试过也有这个现象)不会自动搜索PATH环境变量来加载合适的DLL，而C++是可以的，所以用C++调用的时候一切正常，而Python3.8里面产生了ImportError。
3. python默认加载的是cudart64_110.dll，又不会自动搜索环境变量来加载合适的dll，加载fused.pyd的时候自然就找不到cudart64_12.dll，而提示找不到模块了。
4. 由此可见，编译其实是没有问题的，编译的二进制文件也是可以运行的，只要想办法让python加载正确的dll即可。

解决bug

方案一：更换cuda toolkit版本

安装和pytorch使用的cudart相同版本的cuda toolkit，并配置好环境让ninja用该版本cuda来编译，这样的话，当我们import torch之后，合适版本的cudart会被加载，自然而然就不会提示找不到模块了。如果这个方法依然有问题，可以尝试下一个方法。

方案二：强制python加载合适的dll

我们可以使用类似于以下的代码：

os.add_dll_directory('C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.1\\bin')

这样python会在我们给出的路径里面搜索dll，自然而然就不会找不到dll了。

在测试的时候发现torch.utils.cpp_extension.load在没有代码更改的情况下也运行得很慢，这里暂时没有去深究原因，load代码里面似乎是有版本控制的，通过JIT_EXTENSION_VERSIONER.bump_version_if_changed和JIT_EXTENSION_VERSIONER.get_version的版本比较来决定是否要编译，但是不知道为什么，笔者机器上，虽然源代码没有修改，理论上hash值也没有变化，但是_jit_compile依然会重新编译，不知道是否是Windows的兼容性问题还是其他地方的代码逻辑问题，等下次有时间再仔细阅读一下代码。

为了优化这个问题，笔者封装了一个模块来进行版本控制，如果有类似的需求可以参考:

import os
from torch.utils.cpp_extension import load, _get_build_directory, _import_module_from_library

# 这里根据直接的情况修改路径
os.add_dll_directory('C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.1\\bin')

def getmodule(name, sources, **kargs):
    build_directory = _get_build_directory(name, False)
    module_build_path = os.path.join(
        build_directory, '{}.pyd'.format(name))
    if os.path.exists(module_build_path):
        module_modify_time = os.path.getmtime(module_build_path)
        need_update = False
        for source_file in sources:
            source_modify_time = os.path.getmtime(source_file)
            if source_modify_time > module_modify_time:
                need_update = True
                break
        if not need_update:
            try:
                return _import_module_from_library(name, build_directory, True)
            except:
                pass
    return load(name, sources, **kargs)

这里的原理就是通过比较源文件和二进制文件的修改时间来判断是否要重新编译，这应该是比较快并且比较简单的版本控制的方法。

其他方法

其他博主使用python代码来代替C++和cuda代码，也算是殊途同归，同样也是解决问题，只不过用的不同的思路，笔者比较喜欢去研究问题的成因，然后根据原因来解决。比如在笔者遇到的问题中，发现编译的文件本身没有问题，所以就想办法让python加载正确的dll来解决问题。当然有时候我们也可以选择绕开问题，这也是一种不错的处理问题的方法，比如以上的方法都没法解决找不到dll的问题的时候，如果我们还很赶时间，就可以参考之前那位博主的方法用python重写相关的函数。

其他

PyTorch和CUDA对应版本可以查看官网