本文整理汇总了Python中torch.utils.cpp_extension.load方法的典型用法代码示例。如果您正苦于以下问题：Python cpp_extension.load方法的具体用法？Python cpp_extension.load怎么用？Python cpp_extension.load使用的例子？那么, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在类torch.utils.cpp_extension的用法示例。
在下文中一共展示了cpp_extension.load方法的7个代码示例，这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞，您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: _load_C_extensions
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def _load_C_extensions():
this_dir = os.path.dirname(os.path.abspath(__file__))
this_dir = os.path.dirname(this_dir)
this_dir = os.path.join(this_dir, "csrc")
main_file = glob.glob(os.path.join(this_dir, "*.cpp"))
source_cpu = glob.glob(os.path.join(this_dir, "cpu", "*.cpp"))
source_cuda = glob.glob(os.path.join(this_dir, "cuda", "*.cu"))
source = main_file + source_cpu
extra_cflags = []
if torch.cuda.is_available() and CUDA_HOME is not None:
source.extend(source_cuda)
extra_cflags = ["-DWITH_CUDA"]
source = [os.path.join(this_dir, s) for s in source]
extra_include_paths = [this_dir]
return load_ext(
"torchvision",
source,
extra_cflags=extra_cflags,
extra_include_paths=extra_include_paths,
) 示例2: _load_C_extensions
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def _load_C_extensions():
this_dir = os.path.dirname(os.path.abspath(__file__))
this_dir = os.path.join(this_dir, "csrc")
main_file = glob.glob(os.path.join(this_dir, "*.cpp"))
sources_cpu = glob.glob(os.path.join(this_dir, "cpu", "*.cpp"))
sources_cuda = glob.glob(os.path.join(this_dir, "cuda", "*.cu"))
sources = main_file + sources_cpu
extra_cflags = []
extra_cuda_cflags = []
if torch.cuda.is_available() and CUDA_HOME is not None:
sources.extend(sources_cuda)
extra_cflags = ["-O3", "-DWITH_CUDA"]
extra_cuda_cflags = ["--expt-extended-lambda"]
sources = [os.path.join(this_dir, s) for s in sources]
extra_include_paths = [this_dir]
return load(
name="ext_lib",
sources=sources,
extra_cflags=extra_cflags,
extra_include_paths=extra_include_paths,
extra_cuda_cflags=extra_cuda_cflags,
) 示例3: __init__
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def __init__(self, shapes, verbose=True):
"""
Parameters
----------
shapes : np.ndarray
eigen shapes (obtained by PCA)
"""
super().__init__()
self.register_buffer("_shape_mean",
torch.from_numpy(shapes[0]).float().unsqueeze(0))
components = []
for i, _shape in enumerate(shapes[1:]):
components.append(torch.from_numpy(_shape).float().unsqueeze(0))
component_tensor = torch.cat(components).unsqueeze(0)
self.register_buffer("_shape_components", component_tensor)
self._func = load_cpp("shape_function",
sources=[os.path.join(os.path.split(__file__)[0],
"shape_layer.cpp")],
verbose=verbose) 示例4: load_cpp_ext
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def load_cpp_ext(ext_name):
root_dir = os.path.join(os.path.split(__file__)[0])
src_dir = os.path.join(root_dir, "cpp")
tar_dir = os.path.join(src_dir, "build", ext_name)
os.makedirs(tar_dir, exist_ok=True)
srcs = glob(f"{src_dir}/*.cu") + glob(f"{src_dir}/*.cpp")
with warnings.catch_warnings():
warnings.simplefilter("ignore")
from torch.utils.cpp_extension import load
ext = load(
name=ext_name,
sources=srcs,
extra_cflags=["-O3"],
extra_cuda_cflags=[],
build_directory=tar_dir,
)
return ext
# defer calling load_cpp_ext to make CUDA_VISIBLE_DEVICES happy 示例5: _lazy_load_cpu_kernel
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def _lazy_load_cpu_kernel():
global SRU_CPU_kernel
if SRU_CPU_kernel is not None:
return SRU_CPU_kernel
try:
from torch.utils.cpp_extension import load
cpu_source = os.path.join(os.path.dirname(__file__), "sru_cpu_impl.cpp")
SRU_CPU_kernel = load(
name="sru_cpu_impl",
sources=[cpu_source],
extra_cflags=['-O3'],
verbose=False
)
except:
# use Python version instead
SRU_CPU_kernel = False
return SRU_CPU_kernel
# load C++ implementation for GPU computation 示例6: __init__
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def __init__(self, n_dims, verbose=True):
"""
Parameters
----------
n_dims : int
number of dimensions
verbose : float
if True: verbosity during C++ loading
"""
super().__init__()
homogen_trafo = torch.zeros(1, n_dims + 1, n_dims + 1)
homogen_trafo[:, -1, :-1] = 0.
homogen_trafo[:, -1, -1] = 1.
self.register_buffer("_trafo_matrix", homogen_trafo)
self._n_dims = n_dims
self._func = load_cpp("homogeneous_transform_function",
sources=[
os.path.join(
os.path.split(__file__)[0],
"homogeneous_transform_layer.cpp")],
verbose=verbose) 示例7: _load_C_extensions
# 需要导入模块: from torch.utils import cpp_extension [as 别名]
# 或者: from torch.utils.cpp_extension import load [as 别名]
def _load_C_extensions():
this_dir = os.path.dirname(os.path.abspath(__file__))
this_dir = os.path.dirname(this_dir)
this_dir = os.path.join(this_dir, "csrc")
main_file = glob.glob(os.path.join(this_dir, "*.cpp"))
source_cpu = glob.glob(os.path.join(this_dir, "cpu", "*.cpp"))
source_cuda = glob.glob(os.path.join(this_dir, "cuda", "*.cu"))
source = main_file + source_cpu
extra_cflags = []
if torch.cuda.is_available() and CUDA_HOME is not None:
source.extend(source_cuda)
extra_cflags = ["-DWITH_CUDA"]
source = [os.path.join(this_dir, s) for s in source]
extra_include_paths = [this_dir]
build_dir = os.path.dirname(os.path.abspath(__file__))
build_dir = os.path.dirname(build_dir)
build_dir = os.path.dirname(build_dir)
build_dir = os.path.join(build_dir, 'jit', 'torchvision')
if not os.path.exists(build_dir):
os.makedirs(build_dir)
return load_ext(
"torchvision",
source,
extra_cflags=extra_cflags,
extra_include_paths=extra_include_paths,
build_directory=build_dir,
verbose=True,
)