mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Attempts to fix #92656 BC-breaking! This changes the default of zero_grad in optim and in nn to default set grads to None instead of zero tensors. We are changing the default because there are proven perf wins and existing code has typically not regressed due to this change. (will probably have to flesh out this note more). Pull Request resolved: https://github.com/pytorch/pytorch/pull/92731 Approved by: https://github.com/ngimel
875 lines
32 KiB
Python
875 lines
32 KiB
Python
# Owner(s): ["module: cpp-extensions"]
|
|
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import unittest
|
|
import warnings
|
|
import re
|
|
import tempfile
|
|
import subprocess
|
|
import glob
|
|
|
|
import torch.testing._internal.common_utils as common
|
|
import torch
|
|
import torch.backends.cudnn
|
|
import torch.utils.cpp_extension
|
|
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
|
|
from torch.testing._internal.common_utils import gradcheck
|
|
|
|
|
|
TEST_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
|
|
TEST_CUDNN = False
|
|
TEST_ROCM = torch.cuda.is_available() and torch.version.hip is not None and ROCM_HOME is not None
|
|
if TEST_CUDA and torch.version.cuda is not None: # the skip CUDNN test for ROCm
|
|
CUDNN_HEADER_EXISTS = os.path.isfile(os.path.join(CUDA_HOME, "include/cudnn.h"))
|
|
TEST_CUDNN = (
|
|
TEST_CUDA and CUDNN_HEADER_EXISTS and torch.backends.cudnn.is_available()
|
|
)
|
|
IS_WINDOWS = sys.platform == "win32"
|
|
|
|
|
|
def remove_build_path():
|
|
if sys.platform == "win32":
|
|
print("Not wiping extensions build folder because Windows")
|
|
return
|
|
default_build_root = torch.utils.cpp_extension.get_default_build_root()
|
|
if os.path.exists(default_build_root):
|
|
shutil.rmtree(default_build_root)
|
|
|
|
# There's only one test that runs gracheck, run slow mode manually
|
|
class TestCppExtensionJIT(common.TestCase):
|
|
"""Tests just-in-time cpp extensions.
|
|
Don't confuse this with the PyTorch JIT (aka TorchScript).
|
|
"""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
# cpp extensions use relative paths. Those paths are relative to
|
|
# this file, so we'll change the working directory temporarily
|
|
self.old_working_dir = os.getcwd()
|
|
os.chdir(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
def tearDown(self):
|
|
super().tearDown()
|
|
# return the working directory (see setUp)
|
|
os.chdir(self.old_working_dir)
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
remove_build_path()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
remove_build_path()
|
|
|
|
def test_jit_compile_extension(self):
|
|
module = torch.utils.cpp_extension.load(
|
|
name="jit_extension",
|
|
sources=[
|
|
"cpp_extensions/jit_extension.cpp",
|
|
"cpp_extensions/jit_extension2.cpp",
|
|
],
|
|
extra_include_paths=["cpp_extensions"],
|
|
extra_cflags=["-g"],
|
|
verbose=True,
|
|
)
|
|
x = torch.randn(4, 4)
|
|
y = torch.randn(4, 4)
|
|
|
|
z = module.tanh_add(x, y)
|
|
self.assertEqual(z, x.tanh() + y.tanh())
|
|
|
|
# Checking we can call a method defined not in the main C++ file.
|
|
z = module.exp_add(x, y)
|
|
self.assertEqual(z, x.exp() + y.exp())
|
|
|
|
# Checking we can use this JIT-compiled class.
|
|
doubler = module.Doubler(2, 2)
|
|
self.assertIsNone(doubler.get().grad)
|
|
self.assertEqual(doubler.get().sum(), 4)
|
|
self.assertEqual(doubler.forward().sum(), 8)
|
|
|
|
@unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
|
|
def test_jit_cuda_extension(self):
|
|
# NOTE: The name of the extension must equal the name of the module.
|
|
module = torch.utils.cpp_extension.load(
|
|
name="torch_test_cuda_extension",
|
|
sources=[
|
|
"cpp_extensions/cuda_extension.cpp",
|
|
"cpp_extensions/cuda_extension.cu",
|
|
],
|
|
extra_cuda_cflags=["-O2"],
|
|
verbose=True,
|
|
keep_intermediates=False,
|
|
)
|
|
|
|
x = torch.zeros(100, device="cuda", dtype=torch.float32)
|
|
y = torch.zeros(100, device="cuda", dtype=torch.float32)
|
|
|
|
z = module.sigmoid_add(x, y).cpu()
|
|
|
|
# 2 * sigmoid(0) = 2 * 0.5 = 1
|
|
self.assertEqual(z, torch.ones_like(z))
|
|
|
|
def _run_jit_cuda_archflags(self, flags, expected):
|
|
# Compile an extension with given `flags`
|
|
def _check_cuobjdump_output(expected_values, is_ptx=False):
|
|
elf_or_ptx = '--list-ptx' if is_ptx else '--list-elf'
|
|
lib_ext = '.pyd' if IS_WINDOWS else '.so'
|
|
# Note, .extension name may include _v1, _v2, so first find exact name
|
|
ext_filename = glob.glob(os.path.join(temp_dir,
|
|
'cudaext_archflag*' + lib_ext))[0]
|
|
command = ['cuobjdump', elf_or_ptx, ext_filename]
|
|
p = subprocess.Popen(command,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE)
|
|
output, err = p.communicate()
|
|
output = output.decode("ascii")
|
|
err = err.decode("ascii")
|
|
|
|
if not p.returncode == 0 or not err == '':
|
|
raise AssertionError("Flags: {}\nReturncode: {}\nStderr: {}\n"
|
|
"Output: {} ".format(flags, p.returncode,
|
|
err, output))
|
|
|
|
actual_arches = sorted(re.findall(r'sm_\d\d', output))
|
|
expected_arches = sorted(['sm_' + xx for xx in expected_values])
|
|
self.assertEqual(actual_arches, expected_arches,
|
|
msg="Flags: {}, Actual: {}, Expected: {}\n"
|
|
"Stderr: {}\nOutput: {}".format(
|
|
flags, actual_arches, expected_arches,
|
|
err, output))
|
|
|
|
temp_dir = tempfile.mkdtemp()
|
|
old_envvar = os.environ.get('TORCH_CUDA_ARCH_LIST', None)
|
|
try:
|
|
os.environ['TORCH_CUDA_ARCH_LIST'] = flags
|
|
torch.utils.cpp_extension.load(
|
|
name="cudaext_archflags",
|
|
sources=[
|
|
"cpp_extensions/cuda_extension.cpp",
|
|
"cpp_extensions/cuda_extension.cu",
|
|
],
|
|
extra_cuda_cflags=["-O2"],
|
|
verbose=True,
|
|
build_directory=temp_dir,
|
|
)
|
|
|
|
# Expected output for --list-elf:
|
|
# ELF file 1: cudaext_archflags.1.sm_61.cubin
|
|
# ELF file 2: cudaext_archflags.2.sm_52.cubin
|
|
_check_cuobjdump_output(expected[0])
|
|
if expected[1] is not None:
|
|
# Expected output for --list-ptx:
|
|
# PTX file 1: cudaext_archflags.1.sm_61.ptx
|
|
_check_cuobjdump_output(expected[1], is_ptx=True)
|
|
finally:
|
|
if IS_WINDOWS:
|
|
print("Not wiping extensions build folder because Windows")
|
|
else:
|
|
shutil.rmtree(temp_dir)
|
|
|
|
if old_envvar is None:
|
|
os.environ.pop('TORCH_CUDA_ARCH_LIST')
|
|
else:
|
|
os.environ['TORCH_CUDA_ARCH_LIST'] = old_envvar
|
|
|
|
@unittest.skipIf(not TEST_CUDA, "CUDA not found")
|
|
@unittest.skipIf(TEST_ROCM, "disabled on rocm")
|
|
def test_jit_cuda_archflags(self):
|
|
# Test a number of combinations:
|
|
# - the default for the machine we're testing on
|
|
# - Separators, can be ';' (most common) or ' '
|
|
# - Architecture names
|
|
# - With/without '+PTX'
|
|
|
|
n = torch.cuda.device_count()
|
|
capabilities = {torch.cuda.get_device_capability(i) for i in range(n)}
|
|
# expected values is length-2 tuple: (list of ELF, list of PTX)
|
|
# note: there should not be more than one PTX value
|
|
archflags = {
|
|
'': (['{}{}'.format(capability[0], capability[1]) for capability in capabilities], None),
|
|
"Maxwell+Tegra;6.1": (['53', '61'], None),
|
|
"Volta": (['70'], ['70']),
|
|
}
|
|
if int(torch.version.cuda.split('.')[0]) >= 10:
|
|
# CUDA 9 only supports compute capability <= 7.2
|
|
archflags["7.5+PTX"] = (['75'], ['75'])
|
|
archflags["5.0;6.0+PTX;7.0;7.5"] = (['50', '60', '70', '75'], ['60'])
|
|
if int(torch.version.cuda.split('.')[0]) < 12:
|
|
# CUDA 12 drops compute capability < 5.0
|
|
archflags["Pascal 3.5"] = (['35', '60', '61'], None)
|
|
|
|
for flags, expected in archflags.items():
|
|
self._run_jit_cuda_archflags(flags, expected)
|
|
|
|
@unittest.skipIf(not TEST_CUDNN, "CuDNN not found")
|
|
def test_jit_cudnn_extension(self):
|
|
# implementation of CuDNN ReLU
|
|
if IS_WINDOWS:
|
|
extra_ldflags = ["cudnn.lib"]
|
|
else:
|
|
extra_ldflags = ["-lcudnn"]
|
|
module = torch.utils.cpp_extension.load(
|
|
name="torch_test_cudnn_extension",
|
|
sources=["cpp_extensions/cudnn_extension.cpp"],
|
|
extra_ldflags=extra_ldflags,
|
|
verbose=True,
|
|
with_cuda=True,
|
|
)
|
|
|
|
x = torch.randn(100, device="cuda", dtype=torch.float32)
|
|
y = torch.zeros(100, device="cuda", dtype=torch.float32)
|
|
module.cudnn_relu(x, y) # y=relu(x)
|
|
self.assertEqual(torch.nn.functional.relu(x), y)
|
|
with self.assertRaisesRegex(RuntimeError, "same size"):
|
|
y_incorrect = torch.zeros(20, device="cuda", dtype=torch.float32)
|
|
module.cudnn_relu(x, y_incorrect)
|
|
|
|
def test_inline_jit_compile_extension_with_functions_as_list(self):
|
|
cpp_source = """
|
|
torch::Tensor tanh_add(torch::Tensor x, torch::Tensor y) {
|
|
return x.tanh() + y.tanh();
|
|
}
|
|
"""
|
|
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="inline_jit_extension_with_functions_list",
|
|
cpp_sources=cpp_source,
|
|
functions="tanh_add",
|
|
verbose=True,
|
|
)
|
|
|
|
self.assertEqual(module.tanh_add.__doc__.split("\n")[2], "tanh_add")
|
|
|
|
x = torch.randn(4, 4)
|
|
y = torch.randn(4, 4)
|
|
|
|
z = module.tanh_add(x, y)
|
|
self.assertEqual(z, x.tanh() + y.tanh())
|
|
|
|
def test_inline_jit_compile_extension_with_functions_as_dict(self):
|
|
cpp_source = """
|
|
torch::Tensor tanh_add(torch::Tensor x, torch::Tensor y) {
|
|
return x.tanh() + y.tanh();
|
|
}
|
|
"""
|
|
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="inline_jit_extension_with_functions_dict",
|
|
cpp_sources=cpp_source,
|
|
functions={"tanh_add": "Tanh and then sum :D"},
|
|
verbose=True,
|
|
)
|
|
|
|
self.assertEqual(module.tanh_add.__doc__.split("\n")[2], "Tanh and then sum :D")
|
|
|
|
def test_inline_jit_compile_extension_multiple_sources_and_no_functions(self):
|
|
cpp_source1 = """
|
|
torch::Tensor sin_add(torch::Tensor x, torch::Tensor y) {
|
|
return x.sin() + y.sin();
|
|
}
|
|
"""
|
|
|
|
cpp_source2 = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor sin_add(torch::Tensor x, torch::Tensor y);
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("sin_add", &sin_add, "sin(x) + sin(y)");
|
|
}
|
|
"""
|
|
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="inline_jit_extension",
|
|
cpp_sources=[cpp_source1, cpp_source2],
|
|
verbose=True,
|
|
)
|
|
|
|
x = torch.randn(4, 4)
|
|
y = torch.randn(4, 4)
|
|
|
|
z = module.sin_add(x, y)
|
|
self.assertEqual(z, x.sin() + y.sin())
|
|
|
|
@unittest.skip("Temporarily disabled")
|
|
@unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
|
|
def test_inline_jit_compile_extension_cuda(self):
|
|
cuda_source = """
|
|
__global__ void cos_add_kernel(
|
|
const float* __restrict__ x,
|
|
const float* __restrict__ y,
|
|
float* __restrict__ output,
|
|
const int size) {
|
|
const auto index = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (index < size) {
|
|
output[index] = __cosf(x[index]) + __cosf(y[index]);
|
|
}
|
|
}
|
|
|
|
torch::Tensor cos_add(torch::Tensor x, torch::Tensor y) {
|
|
auto output = torch::zeros_like(x);
|
|
const int threads = 1024;
|
|
const int blocks = (output.numel() + threads - 1) / threads;
|
|
cos_add_kernel<<<blocks, threads>>>(x.data<float>(), y.data<float>(), output.data<float>(), output.numel());
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
# Here, the C++ source need only declare the function signature.
|
|
cpp_source = "torch::Tensor cos_add(torch::Tensor x, torch::Tensor y);"
|
|
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="inline_jit_extension_cuda",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["cos_add"],
|
|
verbose=True,
|
|
)
|
|
|
|
self.assertEqual(module.cos_add.__doc__.split("\n")[2], "cos_add")
|
|
|
|
x = torch.randn(4, 4, device="cuda", dtype=torch.float32)
|
|
y = torch.randn(4, 4, device="cuda", dtype=torch.float32)
|
|
|
|
z = module.cos_add(x, y)
|
|
self.assertEqual(z, x.cos() + y.cos())
|
|
|
|
@unittest.skip("Temporarily disabled")
|
|
@unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
|
|
def test_inline_jit_compile_custom_op_cuda(self):
|
|
cuda_source = """
|
|
__global__ void cos_add_kernel(
|
|
const float* __restrict__ x,
|
|
const float* __restrict__ y,
|
|
float* __restrict__ output,
|
|
const int size) {
|
|
const auto index = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (index < size) {
|
|
output[index] = __cosf(x[index]) + __cosf(y[index]);
|
|
}
|
|
}
|
|
|
|
torch::Tensor cos_add(torch::Tensor x, torch::Tensor y) {
|
|
auto output = torch::zeros_like(x);
|
|
const int threads = 1024;
|
|
const int blocks = (output.numel() + threads - 1) / threads;
|
|
cos_add_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), output.data_ptr<float>(), output.numel());
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
# Here, the C++ source need only declare the function signature.
|
|
cpp_source = """
|
|
#include <torch/library.h>
|
|
torch::Tensor cos_add(torch::Tensor x, torch::Tensor y);
|
|
|
|
TORCH_LIBRARY(inline_jit_extension_custom_op_cuda, m) {
|
|
m.def("cos_add", cos_add);
|
|
}
|
|
"""
|
|
|
|
torch.utils.cpp_extension.load_inline(
|
|
name="inline_jit_extension_custom_op_cuda",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
verbose=True,
|
|
is_python_module=False,
|
|
)
|
|
|
|
x = torch.randn(4, 4, device="cuda", dtype=torch.float32)
|
|
y = torch.randn(4, 4, device="cuda", dtype=torch.float32)
|
|
|
|
z = torch.ops.inline_jit_extension_custom_op_cuda.cos_add(x, y)
|
|
self.assertEqual(z, x.cos() + y.cos())
|
|
|
|
def test_inline_jit_compile_extension_throws_when_functions_is_bad(self):
|
|
with self.assertRaises(ValueError):
|
|
torch.utils.cpp_extension.load_inline(
|
|
name="invalid_jit_extension", cpp_sources="", functions=5
|
|
)
|
|
|
|
def test_lenient_flag_handling_in_jit_extensions(self):
|
|
cpp_source = """
|
|
torch::Tensor tanh_add(torch::Tensor x, torch::Tensor y) {
|
|
return x.tanh() + y.tanh();
|
|
}
|
|
"""
|
|
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="lenient_flag_handling_extension",
|
|
cpp_sources=cpp_source,
|
|
functions="tanh_add",
|
|
extra_cflags=["-g\n\n", "-O0 -Wall"],
|
|
extra_include_paths=[" cpp_extensions\n"],
|
|
verbose=True,
|
|
)
|
|
|
|
x = torch.zeros(100, dtype=torch.float32)
|
|
y = torch.zeros(100, dtype=torch.float32)
|
|
z = module.tanh_add(x, y).cpu()
|
|
self.assertEqual(z, x.tanh() + y.tanh())
|
|
|
|
@unittest.skip("Temporarily disabled")
|
|
@unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
|
|
def test_half_support(self):
|
|
"""
|
|
Checks for an issue with operator< ambiguity for half when certain
|
|
THC headers are included.
|
|
|
|
See https://github.com/pytorch/pytorch/pull/10301#issuecomment-416773333
|
|
for the corresponding issue.
|
|
"""
|
|
cuda_source = """
|
|
template<typename T, typename U>
|
|
__global__ void half_test_kernel(const T* input, U* output) {
|
|
if (input[0] < input[1] || input[0] >= input[1]) {
|
|
output[0] = 123;
|
|
}
|
|
}
|
|
|
|
torch::Tensor half_test(torch::Tensor input) {
|
|
auto output = torch::empty(1, input.options().dtype(torch::kFloat));
|
|
AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "half_test", [&] {
|
|
half_test_kernel<scalar_t><<<1, 1>>>(
|
|
input.data<scalar_t>(),
|
|
output.data<float>());
|
|
});
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="half_test_extension",
|
|
cpp_sources="torch::Tensor half_test(torch::Tensor input);",
|
|
cuda_sources=cuda_source,
|
|
functions=["half_test"],
|
|
verbose=True,
|
|
)
|
|
|
|
x = torch.randn(3, device="cuda", dtype=torch.half)
|
|
result = module.half_test(x)
|
|
self.assertEqual(result[0], 123)
|
|
|
|
def test_reload_jit_extension(self):
|
|
def compile(code):
|
|
return torch.utils.cpp_extension.load_inline(
|
|
name="reloaded_jit_extension",
|
|
cpp_sources=code,
|
|
functions="f",
|
|
verbose=True,
|
|
)
|
|
|
|
module = compile("int f() { return 123; }")
|
|
self.assertEqual(module.f(), 123)
|
|
|
|
module = compile("int f() { return 456; }")
|
|
self.assertEqual(module.f(), 456)
|
|
module = compile("int f() { return 456; }")
|
|
self.assertEqual(module.f(), 456)
|
|
|
|
module = compile("int f() { return 789; }")
|
|
self.assertEqual(module.f(), 789)
|
|
|
|
def test_cpp_frontend_module_has_same_output_as_python(self, dtype=torch.double):
|
|
extension = torch.utils.cpp_extension.load(
|
|
name="cpp_frontend_extension",
|
|
sources="cpp_extensions/cpp_frontend_extension.cpp",
|
|
verbose=True,
|
|
)
|
|
|
|
input = torch.randn(2, 5, dtype=dtype)
|
|
cpp_linear = extension.Net(5, 2)
|
|
cpp_linear.to(dtype)
|
|
python_linear = torch.nn.Linear(5, 2).to(dtype)
|
|
|
|
# First make sure they have the same parameters
|
|
cpp_parameters = dict(cpp_linear.named_parameters())
|
|
with torch.no_grad():
|
|
python_linear.weight.copy_(cpp_parameters["fc.weight"])
|
|
python_linear.bias.copy_(cpp_parameters["fc.bias"])
|
|
|
|
cpp_output = cpp_linear.forward(input)
|
|
python_output = python_linear(input)
|
|
self.assertEqual(cpp_output, python_output)
|
|
|
|
cpp_output.sum().backward()
|
|
python_output.sum().backward()
|
|
|
|
for p in cpp_linear.parameters():
|
|
self.assertFalse(p.grad is None)
|
|
|
|
self.assertEqual(cpp_parameters["fc.weight"].grad, python_linear.weight.grad)
|
|
self.assertEqual(cpp_parameters["fc.bias"].grad, python_linear.bias.grad)
|
|
|
|
def test_cpp_frontend_module_python_inter_op(self):
|
|
extension = torch.utils.cpp_extension.load(
|
|
name="cpp_frontend_extension",
|
|
sources="cpp_extensions/cpp_frontend_extension.cpp",
|
|
verbose=True,
|
|
)
|
|
|
|
# Create a torch.nn.Module which uses the C++ module as a submodule.
|
|
class M(torch.nn.Module):
|
|
def __init__(self):
|
|
super(M, self).__init__()
|
|
self.x = torch.nn.Parameter(torch.tensor(1.0))
|
|
self.net = extension.Net(3, 5)
|
|
|
|
def forward(self, input):
|
|
return self.net.forward(input) + self.x
|
|
|
|
net = extension.Net(5, 2)
|
|
net.double()
|
|
net.to(torch.get_default_dtype())
|
|
self.assertEqual(str(net), "Net")
|
|
|
|
# Further embed the torch.nn.Module into a Sequential, and also add the
|
|
# C++ module as an element of the Sequential.
|
|
sequential = torch.nn.Sequential(M(), torch.nn.Tanh(), net, torch.nn.Sigmoid())
|
|
|
|
input = torch.randn(2, 3)
|
|
# Try calling the module!
|
|
output = sequential.forward(input)
|
|
# The call operator is bound to forward too.
|
|
self.assertEqual(output, sequential(input))
|
|
self.assertEqual(list(output.shape), [2, 2])
|
|
|
|
# Do changes on the module hierarchy.
|
|
old_dtype = torch.get_default_dtype()
|
|
sequential.to(torch.float64)
|
|
sequential.to(torch.float32)
|
|
sequential.to(old_dtype)
|
|
self.assertEqual(sequential[2].parameters()[0].dtype, old_dtype)
|
|
|
|
# Make sure we can access these methods recursively.
|
|
self.assertEqual(len(list(sequential.parameters())), len(net.parameters()) * 2 + 1)
|
|
self.assertEqual(len(list(sequential.named_parameters())), len(net.named_parameters()) * 2 + 1)
|
|
self.assertEqual(len(list(sequential.buffers())), len(net.buffers()) * 2)
|
|
self.assertEqual(len(list(sequential.modules())), 8)
|
|
|
|
# Test clone()
|
|
net2 = net.clone()
|
|
self.assertEqual(len(net.parameters()), len(net2.parameters()))
|
|
self.assertEqual(len(net.buffers()), len(net2.buffers()))
|
|
self.assertEqual(len(net.modules()), len(net2.modules()))
|
|
|
|
# Try differentiating through the whole module.
|
|
for parameter in net.parameters():
|
|
self.assertIsNone(parameter.grad)
|
|
output.sum().backward()
|
|
for parameter in net.parameters():
|
|
self.assertFalse(parameter.grad is None)
|
|
self.assertGreater(parameter.grad.sum(), 0)
|
|
|
|
# Try calling zero_grad()
|
|
net.zero_grad()
|
|
for p in net.parameters():
|
|
assert p.grad is None, "zero_grad defaults to setting grads to None"
|
|
|
|
# Test train(), eval(), training (a property)
|
|
self.assertTrue(net.training)
|
|
net.eval()
|
|
self.assertFalse(net.training)
|
|
net.train()
|
|
self.assertTrue(net.training)
|
|
net.eval()
|
|
|
|
# Try calling the additional methods we registered.
|
|
biased_input = torch.randn(4, 5)
|
|
output_before = net.forward(biased_input)
|
|
bias = net.get_bias().clone()
|
|
self.assertEqual(list(bias.shape), [2])
|
|
net.set_bias(bias + 1)
|
|
self.assertEqual(net.get_bias(), bias + 1)
|
|
output_after = net.forward(biased_input)
|
|
|
|
self.assertNotEqual(output_before, output_after)
|
|
|
|
# Try accessing parameters
|
|
self.assertEqual(len(net.parameters()), 2)
|
|
np = net.named_parameters()
|
|
self.assertEqual(len(np), 2)
|
|
self.assertIn("fc.weight", np)
|
|
self.assertIn("fc.bias", np)
|
|
|
|
self.assertEqual(len(net.buffers()), 1)
|
|
nb = net.named_buffers()
|
|
self.assertEqual(len(nb), 1)
|
|
self.assertIn("buf", nb)
|
|
self.assertEqual(nb[0][1], torch.eye(5))
|
|
|
|
def test_cpp_frontend_module_has_up_to_date_attributes(self):
|
|
extension = torch.utils.cpp_extension.load(
|
|
name="cpp_frontend_extension",
|
|
sources="cpp_extensions/cpp_frontend_extension.cpp",
|
|
verbose=True,
|
|
)
|
|
|
|
net = extension.Net(5, 2)
|
|
|
|
self.assertEqual(len(net._parameters), 0)
|
|
net.add_new_parameter("foo", torch.eye(5))
|
|
self.assertEqual(len(net._parameters), 1)
|
|
|
|
self.assertEqual(len(net._buffers), 1)
|
|
net.add_new_buffer("bar", torch.eye(5))
|
|
self.assertEqual(len(net._buffers), 2)
|
|
|
|
self.assertEqual(len(net._modules), 1)
|
|
net.add_new_submodule("fc2")
|
|
self.assertEqual(len(net._modules), 2)
|
|
|
|
@unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
|
|
def test_cpp_frontend_module_python_inter_op_with_cuda(self):
|
|
extension = torch.utils.cpp_extension.load(
|
|
name="cpp_frontend_extension",
|
|
sources="cpp_extensions/cpp_frontend_extension.cpp",
|
|
verbose=True,
|
|
)
|
|
|
|
net = extension.Net(5, 2)
|
|
for p in net.parameters():
|
|
self.assertTrue(p.device.type == "cpu")
|
|
cpu_parameters = [p.clone() for p in net.parameters()]
|
|
|
|
device = torch.device("cuda", 0)
|
|
net.to(device)
|
|
|
|
for i, p in enumerate(net.parameters()):
|
|
self.assertTrue(p.device.type == "cuda")
|
|
self.assertTrue(p.device.index == 0)
|
|
self.assertEqual(cpu_parameters[i], p)
|
|
|
|
net.cpu()
|
|
net.add_new_parameter("a", torch.eye(5))
|
|
net.add_new_parameter("b", torch.eye(5))
|
|
net.add_new_buffer("c", torch.eye(5))
|
|
net.add_new_buffer("d", torch.eye(5))
|
|
net.add_new_submodule("fc2")
|
|
net.add_new_submodule("fc3")
|
|
|
|
for p in net.parameters():
|
|
self.assertTrue(p.device.type == "cpu")
|
|
|
|
net.cuda()
|
|
|
|
for p in net.parameters():
|
|
self.assertTrue(p.device.type == "cuda")
|
|
|
|
def test_returns_shared_library_path_when_is_python_module_is_true(self):
|
|
source = """
|
|
#include <torch/script.h>
|
|
torch::Tensor func(torch::Tensor x) { return x; }
|
|
static torch::RegisterOperators r("test::func", &func);
|
|
"""
|
|
torch.utils.cpp_extension.load_inline(
|
|
name="is_python_module",
|
|
cpp_sources=source,
|
|
functions="func",
|
|
verbose=True,
|
|
is_python_module=False,
|
|
)
|
|
self.assertEqual(torch.ops.test.func(torch.eye(5)), torch.eye(5))
|
|
|
|
def test_set_default_type_also_changes_aten_default_type(self):
|
|
module = torch.utils.cpp_extension.load_inline(
|
|
name="test_set_default_type",
|
|
cpp_sources="torch::Tensor get() { return torch::empty({}); }",
|
|
functions="get",
|
|
verbose=True,
|
|
)
|
|
|
|
initial_default = torch.get_default_dtype()
|
|
try:
|
|
self.assertEqual(module.get().dtype, initial_default)
|
|
torch.set_default_dtype(torch.float64)
|
|
self.assertEqual(module.get().dtype, torch.float64)
|
|
torch.set_default_dtype(torch.float32)
|
|
self.assertEqual(module.get().dtype, torch.float32)
|
|
torch.set_default_dtype(torch.float16)
|
|
self.assertEqual(module.get().dtype, torch.float16)
|
|
finally:
|
|
torch.set_default_dtype(initial_default)
|
|
|
|
def test_compilation_error_formatting(self):
|
|
# Test that the missing-semicolon error message has linebreaks in it.
|
|
# This'll fail if the message has been munged into a single line.
|
|
# It's hard to write anything more specific as every compiler has it's own
|
|
# error formatting.
|
|
with self.assertRaises(RuntimeError) as e:
|
|
torch.utils.cpp_extension.load_inline(
|
|
name="test_compilation_error_formatting",
|
|
cpp_sources="int main() { return 0 }")
|
|
pattern = r'.*(\\n|\\r).*'
|
|
self.assertNotRegex(str(e), pattern)
|
|
|
|
def test_warning(self):
|
|
# Note: the module created from this source will include the py::key_error
|
|
# symbol. But because of visibility and the fact that it lives in a
|
|
# different compilation unit than pybind, this trips up ubsan even though
|
|
# it is fine. "ubsan.supp" thus needs to contain "vptr:warn_mod.so".
|
|
source = '''
|
|
// error_type:
|
|
// 0: no error
|
|
// 1: torch::TypeError
|
|
// 2: python_error()
|
|
// 3: py::error_already_set
|
|
at::Tensor foo(at::Tensor x, int error_type) {
|
|
std::ostringstream err_stream;
|
|
err_stream << "Error with " << x.type();
|
|
|
|
TORCH_WARN(err_stream.str());
|
|
if(error_type == 1) {
|
|
throw torch::TypeError(err_stream.str().c_str());
|
|
}
|
|
if(error_type == 2) {
|
|
PyObject* obj = PyTuple_New(-1);
|
|
TORCH_CHECK(!obj);
|
|
// Pretend it was caught in a different thread and restored here
|
|
auto e = python_error();
|
|
e.persist();
|
|
e.restore();
|
|
throw e;
|
|
}
|
|
if(error_type == 3) {
|
|
throw py::key_error(err_stream.str());
|
|
}
|
|
return x.cos();
|
|
}
|
|
'''
|
|
|
|
# Ensure double type for hard-coded c name below
|
|
t = torch.rand(2).double()
|
|
cpp_tensor_name = r"CPUDoubleType"
|
|
|
|
# Without error handling, the warnings cannot be catched
|
|
warn_mod = torch.utils.cpp_extension.load_inline(name='warn_mod',
|
|
cpp_sources=[source],
|
|
functions=['foo'],
|
|
with_pytorch_error_handling=False)
|
|
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warn_mod.foo(t, 0)
|
|
self.assertEqual(len(w), 0)
|
|
|
|
with self.assertRaisesRegex(TypeError, t.type()):
|
|
warn_mod.foo(t, 1)
|
|
self.assertEqual(len(w), 0)
|
|
|
|
with self.assertRaisesRegex(SystemError, "bad argument to internal function"):
|
|
warn_mod.foo(t, 2)
|
|
self.assertEqual(len(w), 0)
|
|
|
|
with self.assertRaisesRegex(KeyError, cpp_tensor_name):
|
|
warn_mod.foo(t, 3)
|
|
self.assertEqual(len(w), 0)
|
|
|
|
|
|
warn_mod = torch.utils.cpp_extension.load_inline(name='warn_mod',
|
|
cpp_sources=[source],
|
|
functions=['foo'],
|
|
with_pytorch_error_handling=True)
|
|
|
|
|
|
with warnings.catch_warnings(record=True) as w:
|
|
# Catched with no error should be detected
|
|
warn_mod.foo(t, 0)
|
|
self.assertEqual(len(w), 1)
|
|
|
|
# Catched with cpp error should also be detected
|
|
with self.assertRaisesRegex(TypeError, t.type()):
|
|
warn_mod.foo(t, 1)
|
|
self.assertEqual(len(w), 2)
|
|
|
|
# Catched with python error should also be detected
|
|
with self.assertRaisesRegex(SystemError, "bad argument to internal function"):
|
|
warn_mod.foo(t, 2)
|
|
self.assertEqual(len(w), 3)
|
|
|
|
# Catched with pybind error should also be detected
|
|
# Note that there is no type name translation for pybind errors
|
|
with self.assertRaisesRegex(KeyError, cpp_tensor_name):
|
|
warn_mod.foo(t, 3)
|
|
self.assertEqual(len(w), 4)
|
|
|
|
# Make sure raising warnings are handled properly
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("error")
|
|
|
|
# No error, the warning should raise
|
|
with self.assertRaisesRegex(UserWarning, t.type()):
|
|
warn_mod.foo(t, 0)
|
|
self.assertEqual(len(w), 0)
|
|
|
|
# Another error happened, the warning is ignored
|
|
with self.assertRaisesRegex(TypeError, t.type()):
|
|
warn_mod.foo(t, 1)
|
|
self.assertEqual(len(w), 0)
|
|
|
|
def test_autograd_from_cpp(self):
|
|
source = '''
|
|
void run_back(at::Tensor x) {
|
|
x.backward({});
|
|
}
|
|
|
|
void run_back_no_gil(at::Tensor x) {
|
|
pybind11::gil_scoped_release no_gil;
|
|
x.backward({});
|
|
}
|
|
'''
|
|
|
|
class MyFn(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, x):
|
|
return x.clone()
|
|
|
|
@staticmethod
|
|
def backward(ctx, gx):
|
|
return gx
|
|
|
|
test_backward_deadlock = torch.utils.cpp_extension.load_inline(name='test_backward_deadlock',
|
|
cpp_sources=[source],
|
|
functions=['run_back', 'run_back_no_gil'],)
|
|
|
|
# This used to deadlock
|
|
inp = torch.rand(20, requires_grad=True)
|
|
loss = MyFn.apply(inp).sum()
|
|
with self.assertRaisesRegex(RuntimeError, "The autograd engine was called while holding the GIL."):
|
|
test_backward_deadlock.run_back(loss)
|
|
|
|
inp = torch.rand(20, requires_grad=True)
|
|
loss = MyFn.apply(inp).sum()
|
|
test_backward_deadlock.run_back_no_gil(loss)
|
|
|
|
def test_custom_compound_op_autograd(self):
|
|
# Test that a custom compound op (i.e. a custom op that just calls other aten ops)
|
|
# correctly returns gradients of those other ops
|
|
|
|
source = """
|
|
#include <torch/library.h>
|
|
torch::Tensor my_add(torch::Tensor x, torch::Tensor y) {
|
|
return x + y;
|
|
}
|
|
TORCH_LIBRARY(my, m) {
|
|
m.def("add", &my_add);
|
|
}
|
|
"""
|
|
|
|
torch.utils.cpp_extension.load_inline(
|
|
name="is_python_module",
|
|
cpp_sources=source,
|
|
verbose=True,
|
|
is_python_module=False,
|
|
)
|
|
|
|
a = torch.randn(5, 5, requires_grad=True)
|
|
b = torch.randn(5, 5, requires_grad=True)
|
|
|
|
for fast_mode in (True, False):
|
|
gradcheck(torch.ops.my.add, [a, b], eps=1e-2, fast_mode=fast_mode)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
common.run_tests()
|