PyTorch之自定义op扩展编译与加载 —— cpp_extension

前言

在使用PyTorch或者TensorFlow构建网络时，可能会遇到一些不支持的自定义算子或者对原生算子做一些加速工作，这个时候就需要使用深度学习框架提供的扩展op功能。最典型的便是StyleGAN2中自定义的用于加速和节省内存占用的两个算子。

PyTorch中的op扩展

以PyTorch框架为例，其支持扩展op，支持的函数均在TORCH.UTILS.CPP_EXTENSION中。

load

使用load()方法实时（just-in-time (JIT)）生成op扩展并加载：

torch.utils.cpp_extension.load()

其中编译和生成工具主要使用Ninja，因此需要先安装ninja（pip install ninja）。

运行时会先生成build.ninja文件，其中定义了C++扩展文件（如fused_bias_act.o）和CUDA扩展文件（主要是核函数的实现，如fused_bias_act_kernel.cuda.o）的编译，并定义了扩展文件之间的链接方法。随后执行ninja编译链接过程，最终生成可执行文件（如fused.so）供python脚本调用。

op生成过程中~/.cache/torch_extensions/fused/目录下的文件：

build.ninja  fused_bias_act_kernel.cuda.o  fused_bias_act.o  fused.so

_import_module_from_library

使用_import_module_from_library方法可以加载已经生成好的op扩展（例如fused.so）。

注：其他方法可以查阅 PyTorch Docs > torch.utils.cpp_extension

op扩展生成过程中的问题

Ninja is required to load C++ extensions

Traceback (most recent call last):
  File "python3.8/site-packages/torch/utils/cpp_extension.py", line 752, in verify_ninja_availability
	raise RuntimeError("Ninja is required to load C++ extensions")
RuntimeError: Ninja is required to load C++ extension

使用torch.utils.cpp_extension.load()方法生成op扩展时，其中编译和生成工具主要使用Ninja，因此需要先安装ninja。

pip install ninja

baton.wait()

使用torch.utils.cpp_extension.load()方法生成op扩展时，偶尔会遇到程序一直卡住的问题。此时如果按Ctrl+C取消运行，则会出现类似下面的报错信息：

Traceback (most recent call last):
  File "test.py", line 5, in <module>
    "./fused_bias_act_kernel.cu",
  File "./env/lib/python3.7/site-packages/torch/utils/cpp_extension.py", line 87, in load
    keep_intermediates=keep_intermediates)
  File "./env/lib/python3.7/site-packages/torch/utils/cpp_extension.py", line 997, in _jit_compile
    baton.wait()
  File "./env/lib/python3.7/site-packages/torch/utils/file_baton.py", line 76, in wait
    time.sleep(self.wait_seconds)

从打印信息推测是一直在等待文件锁被释放。此时查看~/.cache/torch_extensions/fused/目录下的文件，可以发现多了一个lock文件，这是一个读文件锁，如果在某次执行编译中杀掉进程而该锁未被释放，则下次再运行编译则会遇到该lock已经存在的问题，导致进程一直等待从而造成假死状态。此时只需先删掉该lock文件后重新执行程序即可。

兼容运行时生成op和直接加载op

当自定义op一旦编译生成好，则下次运行时即可直接加载使用而不用再执行编译过程，这样可以大大节省编译时间。

以fused_bias_act op为例：

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

try:
    user_home_path = os.path.expanduser('~')
    fused = _import_module_from_library('fused', user_home_path+'/.cache/torch_extensions/fused', True)
    print(f'Load fused from {user_home_path}/.cache/torch_extensions/fused')
except:
    module_path = os.path.dirname(__file__)
    fused = load(
        'fused',
        sources=[
            os.path.join(module_path, 'fused_bias_act.cpp'),
            os.path.join(module_path, 'fused_bias_act_kernel.cu'),
        ],
    )
    print(f'Build fused from cpp & cu files')

版权说明

本文为原创文章，独家发布在blog.csdn.net/TracelessLe。未经个人允许不得转载。如需帮助请email至[email protected]。

参考资料

[1] NVlabs/stylegan2: StyleGAN2 - Official TensorFlow Implementation
[2] PyTorch Docs > torch.utils.cpp_extension
[3] python - Ninja is required to load C++ extensions - Stack Overflow
[4] torch.utils.cpp_extension.load卡住无响应_zParquet的博客-CSDN博客
[5] _import_module_from_library