W3Cschool
恭喜您成為首批注冊用戶
獲得88經(jīng)驗值獎勵
PyTorch C++ 擴展開發(fā)詳解
在深度學習開發(fā)中,我們有時會遇到 Python 實現(xiàn)的性能瓶頸,或者需要調(diào)用現(xiàn)有的 C++_nlank 或 CUDA 代碼。PyTorch 提供了強大的 C++ 擴展功能,允許開發(fā)者將自定義的 C++ 或 CUDA 代碼集成到 PyTorch 模型中,從而提升性能或復(fù)用代碼。
torch.utils.cpp_extension.CppExtension
用于創(chuàng)建一個 setuptools.Extension
,用于構(gòu)建 C++ 擴展。
name
:擴展的名稱。sources
:C++ 源文件的路徑列表。setuptools.Extension
構(gòu)造函數(shù)。from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
name='extension',
ext_modules=[
CppExtension(
name='extension',
sources=['extension.cpp'],
extra_compile_args=['-g']
)
],
cmdclass={
'build_ext': BuildExtension
}
)
torch.utils.cpp_extension.CUDAExtension
用于創(chuàng)建一個 setuptools.Extension
,用于構(gòu)建 CUDA/C++ 擴展。它會自動包含 CUDA 的相關(guān)路徑和庫。
name
:擴展的名稱。sources
:C++ 和 CUDA 源文件的路徑列表。setuptools.Extension
構(gòu)造函數(shù)。from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
name='cuda_extension',
ext_modules=[
CUDAExtension(
name='cuda_extension',
sources=['extension.cpp', 'extension_kernel.cu'],
extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']}
)
],
cmdclass={
'build_ext': BuildExtension
}
)
torch.utils.cpp_extension.BuildExtension
自定義 setuptools
構(gòu)建擴展的類。它負責傳遞編譯器標志,并支持混合 C++/CUDA 編譯。
torch.utils.cpp_extension.load
即時加載 PyTorch C++ 擴展(JIT)。它會編譯源代碼并將其作為模塊加載到當前 Python 進程中。
name
:擴展的名稱。sources
:C++ 源文件的路徑列表。extra_cflags
:傳遞給編譯器的額外標志。extra_cuda_cflags
:傳遞給 nvcc 的額外標志。extra_ldflags
:傳遞給鏈接器的額外標志。build_directory
:構(gòu)建工作區(qū)的路徑。verbose
:是否開啟詳細日志。with_cuda
:是否包含 CUDA 支持。is_python_module
:是否作為 Python 模塊加載。from torch.utils.cpp_extension import load
module = load(
name='extension',
sources=['extension.cpp', 'extension_kernel.cu'],
extra_cflags=['-O2'],
verbose=True
)
torch.utils.cpp_extension.load_inline
從字符串源即時加載 PyTorch C++ 擴展(JIT)。它與 load
類似,但源代碼以字符串形式提供。
name
:擴展的名稱。cpp_sources
:C++ 源代碼字符串或字符串列表。cuda_sources
:CUDA 源代碼字符串或字符串列表。functions
:需要生成綁定的函數(shù)名稱列表或字典。load
類似。from torch.utils.cpp_extension import load_inline
source = '''
at::Tensor sin_add(at::Tensor x, at::Tensor y) {
return x.sin() + y.sin();
}
'''
module = load_inline(
name='inline_extension',
cpp_sources=[source],
functions=['sin_add']
)
torch.utils.cpp_extension.include_paths(cuda=False)
:獲取構(gòu)建 C++ 或 CUDA 擴展所需的包含路徑。torch.utils.cpp_extension.check_compiler_abi_compatibility(compiler)
:驗證編譯器是否與 PyTorch 兼容。torch.utils.cpp_extension.verify_ninja_availability()
:檢查系統(tǒng)上是否有 ninja 構(gòu)建系統(tǒng)。編寫 C++ 或 CUDA 源代碼,實現(xiàn)自定義的功能或操作。
創(chuàng)建一個 setup.py
腳本,使用 CppExtension
或 CUDAExtension
定義擴展,并使用 BuildExtension
構(gòu)建擴展。
運行以下命令構(gòu)建和安裝擴展:
python setup.py install
使用 load
或 load_inline
函數(shù)即時加載擴展,無需編寫 setup.py
腳本。
通過本教程,我們詳細介紹了 PyTorch C++ 擴展的開發(fā)工具和流程。利用這些工具,開發(fā)者可以輕松地將自定義的 C++ 或 CUDA 代碼集成到 PyTorch 項目中,從而提升模型性能或復(fù)用現(xiàn)有代碼。掌握 C++ 擴展的開發(fā)方法,能夠幫助您在深度學習項目中更好地優(yōu)化性能和利用硬件資源。
Copyright©2021 w3cschool編程獅|閩ICP備15016281號-3|閩公網(wǎng)安備35020302033924號
違法和不良信息舉報電話:173-0602-2364|舉報郵箱:jubao@eeedong.com
掃描二維碼
下載編程獅App
編程獅公眾號
聯(lián)系方式:
更多建議: