mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[codemod][lint][fbcode/c*] Enable BLACK by default
Test Plan: manual inspection & sandcastle Reviewed By: zertosh Differential Revision: D30279364 fbshipit-source-id: c1ed77dfe43a3bde358f92737cd5535ae5d13c9a
This commit is contained in:
Zsolt Dollenstein
committed by
Facebook GitHub Bot
Facebook GitHub Bot
parent
aac3c7bd06
commit
b004307252
@ -1,10 +1,11 @@
|
||||
import copy
|
||||
import itertools
|
||||
import functools
|
||||
import itertools
|
||||
import unittest
|
||||
|
||||
try:
|
||||
import torchvision
|
||||
|
||||
HAS_TORCHVISION = True
|
||||
except ImportError:
|
||||
HAS_TORCHVISION = False
|
||||
@ -12,12 +13,18 @@ except ImportError:
|
||||
skipIfNoTorchVision = unittest.skipIf(not HAS_TORCHVISION, "no torchvision")
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.jit
|
||||
import torch.backends.mkldnn
|
||||
import torch.jit
|
||||
import torch.nn.functional as F
|
||||
from torch.testing._internal.common_utils import (
|
||||
TestCase,
|
||||
run_tests,
|
||||
TemporaryFileName,
|
||||
gradcheck,
|
||||
gradgradcheck,
|
||||
IS_WINDOWS,
|
||||
)
|
||||
from torch.utils import mkldnn as mkldnn_utils
|
||||
from torch.testing._internal.common_utils import TestCase, \
|
||||
run_tests, TemporaryFileName, gradcheck, gradgradcheck, IS_WINDOWS
|
||||
|
||||
# batched grad doesn't support mkldnn
|
||||
gradcheck = functools.partial(gradcheck, check_batched_grad=False)
|
||||
@ -30,22 +37,26 @@ gradgradcheck = functools.partial(gradgradcheck, check_batched_grad=False)
|
||||
@functools.lru_cache(maxsize=None)
|
||||
def has_bf16_support():
|
||||
import sys
|
||||
if sys.platform != 'linux':
|
||||
|
||||
if sys.platform != "linux":
|
||||
return False
|
||||
with open("/proc/cpuinfo", encoding="ascii") as f:
|
||||
lines = f.read()
|
||||
return all(word in lines for word in ["avx512bw", "avx512vl", "avx512dq"])
|
||||
|
||||
|
||||
types = [torch.float, torch.bfloat16]
|
||||
|
||||
# Comment the line below to find out the CI machines having MKL-DNN build disabled
|
||||
@unittest.skipIf(not torch._C.has_mkldnn, "MKL-DNN build is disabled")
|
||||
class TestMkldnn(TestCase):
|
||||
def test_conversion(self):
|
||||
for cpu_tensor in [torch.randn((1, 2, 3, 4),
|
||||
dtype=torch.float, device=torch.device('cpu')),
|
||||
torch.randn((1, 2, 3, 4, 5),
|
||||
dtype=torch.float, device=torch.device('cpu'))[:, :, :, :, 1]]:
|
||||
for cpu_tensor in [
|
||||
torch.randn((1, 2, 3, 4), dtype=torch.float, device=torch.device("cpu")),
|
||||
torch.randn((1, 2, 3, 4, 5), dtype=torch.float, device=torch.device("cpu"))[
|
||||
:, :, :, :, 1
|
||||
],
|
||||
]:
|
||||
cpu_tensor.requires_grad_()
|
||||
# float cpu tensor to mkldnn float tensor or bfloat tensor.
|
||||
for dtype1 in types:
|
||||
@ -58,19 +69,31 @@ class TestMkldnn(TestCase):
|
||||
for dtype2 in types:
|
||||
cpu_tensor_2 = mkldnn_tensor.to_dense(dtype2)
|
||||
self.assertEqual(cpu_tensor_2.dtype, dtype2)
|
||||
atol = 1e-5 if dtype1 == torch.float and dtype2 == torch.float else 1e-2
|
||||
self.assertEqual(cpu_tensor, cpu_tensor_2.float(), atol=atol, rtol=0)
|
||||
atol = (
|
||||
1e-5
|
||||
if dtype1 == torch.float and dtype2 == torch.float
|
||||
else 1e-2
|
||||
)
|
||||
self.assertEqual(
|
||||
cpu_tensor, cpu_tensor_2.float(), atol=atol, rtol=0
|
||||
)
|
||||
|
||||
self.assertEqual(mkldnn_tensor.device, torch.device('cpu'))
|
||||
self.assertEqual(mkldnn_tensor.device, torch.device("cpu"))
|
||||
self.assertEqual(mkldnn_tensor.size(), torch.Size([1, 2, 3, 4]))
|
||||
self.assertEqual(mkldnn_tensor.numel(), cpu_tensor.numel())
|
||||
if dtype1 == torch.float:
|
||||
self.assertEqual(mkldnn_tensor.element_size(), cpu_tensor.element_size())
|
||||
self.assertEqual(
|
||||
mkldnn_tensor.element_size(), cpu_tensor.element_size()
|
||||
)
|
||||
else:
|
||||
self.assertEqual(mkldnn_tensor.element_size(), cpu_tensor.element_size() / 2)
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
"Cannot access data pointer of Tensor that doesn't have storage",
|
||||
lambda: mkldnn_tensor.data_ptr() != 0)
|
||||
self.assertEqual(
|
||||
mkldnn_tensor.element_size(), cpu_tensor.element_size() / 2
|
||||
)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"Cannot access data pointer of Tensor that doesn't have storage",
|
||||
lambda: mkldnn_tensor.data_ptr() != 0,
|
||||
)
|
||||
|
||||
# bfloat cpu tensor to mkldnn float tensor or bfloat tensor.
|
||||
cpu_tensor_bf16 = cpu_tensor.bfloat16()
|
||||
@ -84,18 +107,26 @@ class TestMkldnn(TestCase):
|
||||
for dtype2 in types:
|
||||
cpu_tensor_2 = mkldnn_tensor.to_dense(dtype2)
|
||||
self.assertEqual(cpu_tensor_2.dtype, dtype2)
|
||||
self.assertEqual(cpu_tensor_bf16, cpu_tensor_2.bfloat16(), atol=1e-5, rtol=0)
|
||||
self.assertEqual(
|
||||
cpu_tensor_bf16, cpu_tensor_2.bfloat16(), atol=1e-5, rtol=0
|
||||
)
|
||||
|
||||
self.assertEqual(mkldnn_tensor.device, torch.device('cpu'))
|
||||
self.assertEqual(mkldnn_tensor.device, torch.device("cpu"))
|
||||
self.assertEqual(mkldnn_tensor.size(), torch.Size([1, 2, 3, 4]))
|
||||
self.assertEqual(mkldnn_tensor.numel(), cpu_tensor.numel())
|
||||
if dtype1 == torch.bfloat16:
|
||||
self.assertEqual(mkldnn_tensor.element_size(), cpu_tensor_bf16.element_size())
|
||||
self.assertEqual(
|
||||
mkldnn_tensor.element_size(), cpu_tensor_bf16.element_size()
|
||||
)
|
||||
else:
|
||||
self.assertEqual(mkldnn_tensor.element_size(), cpu_tensor_bf16.element_size() * 2)
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
"Cannot access data pointer of Tensor that doesn't have storage",
|
||||
lambda: mkldnn_tensor.data_ptr() != 0)
|
||||
self.assertEqual(
|
||||
mkldnn_tensor.element_size(), cpu_tensor_bf16.element_size() * 2
|
||||
)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"Cannot access data pointer of Tensor that doesn't have storage",
|
||||
lambda: mkldnn_tensor.data_ptr() != 0,
|
||||
)
|
||||
|
||||
def test_copy(self):
|
||||
x = torch.randn(4, 5, dtype=torch.float32)
|
||||
@ -104,35 +135,62 @@ class TestMkldnn(TestCase):
|
||||
mkldnn_z = torch.randn(4, 10, dtype=torch.float32).to_mkldnn()
|
||||
mkldnn_y.copy_(mkldnn_x)
|
||||
self.assertEqual(x, mkldnn_y.to_dense())
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
"copy_mkldnn_: only support same size tensor.",
|
||||
lambda: mkldnn_z.copy_(mkldnn_x))
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
"copy_mkldnn_: between mkldnn layout and dense Tensors is not implemented! "
|
||||
"Found self type = torch.FloatTensor and src type = Mkldnntorch.FloatTensor",
|
||||
lambda: x.copy_(mkldnn_x))
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
"copy_mkldnn_: between mkldnn layout and dense Tensors is not implemented! "
|
||||
"Found self type = Mkldnntorch.FloatTensor and src type = torch.FloatTensor",
|
||||
lambda: mkldnn_x.copy_(x))
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"copy_mkldnn_: only support same size tensor.",
|
||||
lambda: mkldnn_z.copy_(mkldnn_x),
|
||||
)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"copy_mkldnn_: between mkldnn layout and dense Tensors is not implemented! "
|
||||
"Found self type = torch.FloatTensor and src type = Mkldnntorch.FloatTensor",
|
||||
lambda: x.copy_(mkldnn_x),
|
||||
)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"copy_mkldnn_: between mkldnn layout and dense Tensors is not implemented! "
|
||||
"Found self type = Mkldnntorch.FloatTensor and src type = torch.FloatTensor",
|
||||
lambda: mkldnn_x.copy_(x),
|
||||
)
|
||||
|
||||
def test_unsupported(self):
|
||||
# unsupported types and unsupported types with gpu
|
||||
for dtype in [torch.double, torch.half, torch.uint8, torch.int8,
|
||||
torch.short, torch.int, torch.long]:
|
||||
for dtype in [
|
||||
torch.double,
|
||||
torch.half,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.short,
|
||||
torch.int,
|
||||
torch.long,
|
||||
]:
|
||||
with self.assertRaises(RuntimeError) as context:
|
||||
torch.randn(1, 2, 3, 4, dtype=dtype, device=torch.device('cpu')).to_mkldnn()
|
||||
torch.randn(
|
||||
1, 2, 3, 4, dtype=dtype, device=torch.device("cpu")
|
||||
).to_mkldnn()
|
||||
if torch.cuda.is_available():
|
||||
with self.assertRaises(RuntimeError) as context:
|
||||
torch.randn(1, 2, 3, 4, dtype=dtype, device=torch.device('cuda')).to_mkldnn()
|
||||
torch.randn(
|
||||
1, 2, 3, 4, dtype=dtype, device=torch.device("cuda")
|
||||
).to_mkldnn()
|
||||
# supported type with gpu
|
||||
if torch.cuda.is_available():
|
||||
with self.assertRaises(RuntimeError) as context:
|
||||
torch.randn(1, 2, 3, 4, dtype=torch.float, device=torch.device('cuda')).to_mkldnn()
|
||||
torch.randn(
|
||||
1, 2, 3, 4, dtype=torch.float, device=torch.device("cuda")
|
||||
).to_mkldnn()
|
||||
# some factory functions
|
||||
for creator in [torch.ones, torch.randn, torch.rand]:
|
||||
with self.assertRaises(RuntimeError) as context:
|
||||
creator(1, 2, 3, 4, dtype=torch.float, device=torch.device('cpu'), layout=torch._mkldnn)
|
||||
creator(
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
4,
|
||||
dtype=torch.float,
|
||||
device=torch.device("cpu"),
|
||||
layout=torch._mkldnn,
|
||||
)
|
||||
|
||||
def test_autograd_to_mkldnn(self):
|
||||
# MKLDNN only supports float32
|
||||
@ -143,12 +201,16 @@ class TestMkldnn(TestCase):
|
||||
|
||||
# because MKLDNN only supports float32, we need to lessen the precision.
|
||||
# these numbers are just empirical results that seem to work.
|
||||
self.assertWarnsRegex(UserWarning,
|
||||
'double precision floating point',
|
||||
lambda: gradcheck(func, [root], atol=4e-2, rtol=1e-2))
|
||||
self.assertWarnsRegex(UserWarning,
|
||||
'double precision floating point',
|
||||
lambda: gradgradcheck(func, [root], atol=4e-2, rtol=1e-2))
|
||||
self.assertWarnsRegex(
|
||||
UserWarning,
|
||||
"double precision floating point",
|
||||
lambda: gradcheck(func, [root], atol=4e-2, rtol=1e-2),
|
||||
)
|
||||
self.assertWarnsRegex(
|
||||
UserWarning,
|
||||
"double precision floating point",
|
||||
lambda: gradgradcheck(func, [root], atol=4e-2, rtol=1e-2),
|
||||
)
|
||||
|
||||
def test_autograd_from_mkldnn(self):
|
||||
# MKLDNN only supports float32
|
||||
@ -159,9 +221,11 @@ class TestMkldnn(TestCase):
|
||||
|
||||
# because MKLDNN only supports float32, we need to lessen the precision.
|
||||
# these numbers are just empirical results that seem to work.
|
||||
self.assertWarnsRegex(UserWarning,
|
||||
'double precision floating point',
|
||||
lambda: gradcheck(func, [root], atol=4e-2, rtol=1e-2))
|
||||
self.assertWarnsRegex(
|
||||
UserWarning,
|
||||
"double precision floating point",
|
||||
lambda: gradcheck(func, [root], atol=4e-2, rtol=1e-2),
|
||||
)
|
||||
|
||||
def test_detach(self):
|
||||
root = torch.randn(4, 5, dtype=torch.float32).to_mkldnn().requires_grad_()
|
||||
@ -177,8 +241,14 @@ class TestMkldnn(TestCase):
|
||||
self.assertFalse(root.requires_grad)
|
||||
|
||||
def test_repr(self):
|
||||
self.assertTrue("layout=torch._mkldnn" in str(torch.randn((1, 2, 3, 4),
|
||||
dtype=torch.float, device=torch.device('cpu')).to_mkldnn()))
|
||||
self.assertTrue(
|
||||
"layout=torch._mkldnn"
|
||||
in str(
|
||||
torch.randn(
|
||||
(1, 2, 3, 4), dtype=torch.float, device=torch.device("cpu")
|
||||
).to_mkldnn()
|
||||
)
|
||||
)
|
||||
|
||||
def _test_conv_base(self, dim):
|
||||
conv_module = {1: torch.nn.Conv1d, 2: torch.nn.Conv2d, 3: torch.nn.Conv3d}
|
||||
@ -190,14 +260,16 @@ class TestMkldnn(TestCase):
|
||||
C = torch.randint(1, 3, (1,)).item() * groups
|
||||
x_shape = (N, C) + input_shapes[dim]
|
||||
x = torch.randn(x_shape, dtype=torch.float32)
|
||||
conv = conv_module[dim](in_channels=C,
|
||||
out_channels=M,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
groups=groups).float()
|
||||
conv = conv_module[dim](
|
||||
in_channels=C,
|
||||
out_channels=M,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
groups=groups,
|
||||
).float()
|
||||
x1 = x.clone()
|
||||
x2 = x.clone().to_mkldnn()
|
||||
if not train:
|
||||
@ -223,10 +295,9 @@ class TestMkldnn(TestCase):
|
||||
loss2.backward()
|
||||
self.assertTrue(x2.grad.is_mkldnn)
|
||||
self.assertEqual(x1.grad, x2.grad.to_dense())
|
||||
self.assertEqual(conv.weight.grad,
|
||||
mkldnn_conv.weight.grad,
|
||||
atol=1e-3,
|
||||
rtol=1e-3)
|
||||
self.assertEqual(
|
||||
conv.weight.grad, mkldnn_conv.weight.grad, atol=1e-3, rtol=1e-3
|
||||
)
|
||||
if bias:
|
||||
self.assertEqual(conv.bias.grad, mkldnn_conv.bias.grad)
|
||||
|
||||
@ -251,26 +322,34 @@ class TestMkldnn(TestCase):
|
||||
x_shape = (N, C) + input_shapes[dim]
|
||||
x = torch.randn(x_shape, dtype=torch.float32)
|
||||
|
||||
conv = conv_module[dim](in_channels=C,
|
||||
out_channels=M,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
groups=groups).float()
|
||||
conv = conv_module[dim](
|
||||
in_channels=C,
|
||||
out_channels=M,
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
groups=groups,
|
||||
).float()
|
||||
x_bf16 = x.bfloat16()
|
||||
if has_bf16_support():
|
||||
mkldnn_conv = mkldnn_utils.to_mkldnn(copy.deepcopy(conv))
|
||||
mkldnn_conv_bf16 = mkldnn_utils.to_mkldnn(copy.deepcopy(conv), torch.bfloat16)
|
||||
mkldnn_conv_bf16 = mkldnn_utils.to_mkldnn(
|
||||
copy.deepcopy(conv), torch.bfloat16
|
||||
)
|
||||
y = mkldnn_conv(x.to_mkldnn()).to_dense()
|
||||
y_bf16 = mkldnn_conv_bf16(x_bf16.to_mkldnn()).to_dense(torch.float32)
|
||||
self.assertEqual(y, y_bf16, atol=1e-1, rtol=1e-3)
|
||||
else:
|
||||
msg = r"bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
with self.assertRaisesRegex(RuntimeError, msg):
|
||||
mkldnn_conv_bf16 = mkldnn_utils.to_mkldnn(copy.deepcopy(conv), torch.bfloat16)
|
||||
y_bf16 = mkldnn_conv_bf16(x_bf16.to_mkldnn()).to_dense(torch.float32)
|
||||
mkldnn_conv_bf16 = mkldnn_utils.to_mkldnn(
|
||||
copy.deepcopy(conv), torch.bfloat16
|
||||
)
|
||||
y_bf16 = mkldnn_conv_bf16(x_bf16.to_mkldnn()).to_dense(
|
||||
torch.float32
|
||||
)
|
||||
|
||||
def test_conv1d_bf16(self):
|
||||
self._test_conv_bf16_base(dim=1)
|
||||
@ -303,9 +382,7 @@ class TestMkldnn(TestCase):
|
||||
|
||||
self.assertEqual(conv2d_mkldnn.weight.ndimension(), 5)
|
||||
self.assertEqual(conv2d_loaded.weight.ndimension(), 4)
|
||||
self.assertEqual(
|
||||
conv2d(x),
|
||||
conv2d_loaded(x.to_mkldnn()).to_dense())
|
||||
self.assertEqual(conv2d(x), conv2d_loaded(x.to_mkldnn()).to_dense())
|
||||
|
||||
def test_relu(self):
|
||||
x = torch.randn((4, 5), dtype=torch.float32) * 10
|
||||
@ -344,9 +421,7 @@ class TestMkldnn(TestCase):
|
||||
self.assertEqual(y, y_bf16, atol=1e-1, rtol=1e-3)
|
||||
else:
|
||||
msg = r"bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: fn(x_bf16.to_mkldnn()))
|
||||
self.assertRaisesRegex(RuntimeError, msg, lambda: fn(x_bf16.to_mkldnn()))
|
||||
|
||||
def test_relu_bf16(self):
|
||||
self._test_relu_bf16_base("relu")
|
||||
@ -382,12 +457,12 @@ class TestMkldnn(TestCase):
|
||||
loss1.backward()
|
||||
loss2.backward()
|
||||
self.assertEqual(y1, y2.to(torch.float32), atol=1e-1, rtol=0)
|
||||
self.assertEqual(x1.grad.to_dense(), x2.grad.to_dense(torch.float32), atol=1e-2, rtol=0)
|
||||
self.assertEqual(
|
||||
x1.grad.to_dense(), x2.grad.to_dense(torch.float32), atol=1e-2, rtol=0
|
||||
)
|
||||
else:
|
||||
msg = r"bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: m(x2))
|
||||
self.assertRaisesRegex(RuntimeError, msg, lambda: m(x2))
|
||||
|
||||
def _test_max_pool_base(self, dim, input):
|
||||
pool_module = {2: torch.nn.MaxPool2d, 3: torch.nn.MaxPool3d}
|
||||
@ -397,7 +472,8 @@ class TestMkldnn(TestCase):
|
||||
kernel_size=3 if not ceil_mode else 7,
|
||||
stride=stride,
|
||||
padding=1,
|
||||
ceil_mode=ceil_mode)
|
||||
ceil_mode=ceil_mode,
|
||||
)
|
||||
|
||||
x1 = input.clone().requires_grad_()
|
||||
x2 = input.clone().to_mkldnn().requires_grad_()
|
||||
@ -424,7 +500,6 @@ class TestMkldnn(TestCase):
|
||||
x = torch.randn(N, C, D, H, W, dtype=torch.float32) * 10
|
||||
self._test_max_pool_base(dim=3, input=x)
|
||||
|
||||
|
||||
@unittest.skipIf(IS_WINDOWS, "Limit support for bf16 path")
|
||||
def _test_max_pool_bf16_base(self, dim, input):
|
||||
pool_module = {2: torch.nn.MaxPool2d, 3: torch.nn.MaxPool3d}
|
||||
@ -435,17 +510,21 @@ class TestMkldnn(TestCase):
|
||||
kernel_size=3 if not ceil_mode else 7,
|
||||
stride=stride,
|
||||
padding=1,
|
||||
ceil_mode=ceil_mode)
|
||||
ceil_mode=ceil_mode,
|
||||
)
|
||||
|
||||
if has_bf16_support():
|
||||
y = max_pool(input.to_mkldnn()).to_dense()
|
||||
y_bf16 = max_pool(x_bf16.to_mkldnn()).to_dense(torch.float32)
|
||||
self.assertEqual(y, y_bf16, atol=0.1, rtol=1e-3)
|
||||
else:
|
||||
msg = "mkldnn_max_pool%dd: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq" % dim
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: max_pool(x_bf16.to_mkldnn()))
|
||||
msg = (
|
||||
"mkldnn_max_pool%dd: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
% dim
|
||||
)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError, msg, lambda: max_pool(x_bf16.to_mkldnn())
|
||||
)
|
||||
|
||||
def test_max_pool2d_bf16(self):
|
||||
N = torch.randint(3, 10, (1,)).item()
|
||||
@ -473,14 +552,16 @@ class TestMkldnn(TestCase):
|
||||
kernel_size=3 if not ceil_mode else 7,
|
||||
stride=None,
|
||||
padding=1,
|
||||
ceil_mode=ceil_mode)
|
||||
ceil_mode=ceil_mode,
|
||||
)
|
||||
|
||||
y2 = F.max_pool2d(
|
||||
x.to_mkldnn(),
|
||||
kernel_size=3 if not ceil_mode else 7,
|
||||
stride=None,
|
||||
padding=1,
|
||||
ceil_mode=ceil_mode)
|
||||
ceil_mode=ceil_mode,
|
||||
)
|
||||
|
||||
self.assertEqual(y1, y2.to_dense())
|
||||
|
||||
@ -491,34 +572,28 @@ class TestMkldnn(TestCase):
|
||||
|
||||
# 2d dilation case
|
||||
x = torch.randn(N, C, 7, 7, dtype=torch.float32).to_mkldnn()
|
||||
max_pool2d = torch.nn.MaxPool2d(
|
||||
kernel_size=3,
|
||||
stride=3,
|
||||
padding=1,
|
||||
dilation=2)
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
'mkldnn_max_pool2d does not support dilation case',
|
||||
lambda: max_pool2d(x))
|
||||
max_pool2d = torch.nn.MaxPool2d(kernel_size=3, stride=3, padding=1, dilation=2)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"mkldnn_max_pool2d does not support dilation case",
|
||||
lambda: max_pool2d(x),
|
||||
)
|
||||
|
||||
# 3d dilation case
|
||||
x = torch.randn(N, C, 7, 7, 7, dtype=torch.float32).to_mkldnn()
|
||||
max_pool3d = torch.nn.MaxPool3d(
|
||||
kernel_size=3,
|
||||
stride=3,
|
||||
padding=1,
|
||||
dilation=2)
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
'mkldnn_max_pool3d does not support dilation case',
|
||||
lambda: max_pool3d(x))
|
||||
max_pool3d = torch.nn.MaxPool3d(kernel_size=3, stride=3, padding=1, dilation=2)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"mkldnn_max_pool3d does not support dilation case",
|
||||
lambda: max_pool3d(x),
|
||||
)
|
||||
|
||||
def _test_avg_pool_base(self, dim, input):
|
||||
avg_module = {2: torch.nn.AvgPool2d, 3: torch.nn.AvgPool3d}
|
||||
for count_include_pad in [True, False]:
|
||||
avg_pool = avg_module[dim](
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
count_include_pad=count_include_pad)
|
||||
kernel_size=3, stride=2, padding=1, count_include_pad=count_include_pad
|
||||
)
|
||||
|
||||
x1 = input.clone().requires_grad_()
|
||||
x2 = input.clone().to_mkldnn().requires_grad_()
|
||||
@ -549,19 +624,20 @@ class TestMkldnn(TestCase):
|
||||
x_bf16 = input.bfloat16()
|
||||
for count_include_pad in [True, False]:
|
||||
avg_pool = avg_module[dim](
|
||||
kernel_size=3,
|
||||
stride=2,
|
||||
padding=1,
|
||||
count_include_pad=count_include_pad)
|
||||
kernel_size=3, stride=2, padding=1, count_include_pad=count_include_pad
|
||||
)
|
||||
if has_bf16_support():
|
||||
y = avg_pool(input.to_mkldnn()).to_dense()
|
||||
y_bf16 = avg_pool(x_bf16.to_mkldnn()).to_dense(torch.float)
|
||||
self.assertEqual(y, y_bf16, atol=1e-1, rtol=1e-3)
|
||||
else:
|
||||
msg = "mkldnn_avg_pool%dd: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq" % dim
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: avg_pool(x_bf16.to_mkldnn()))
|
||||
msg = (
|
||||
"mkldnn_avg_pool%dd: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
% dim
|
||||
)
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError, msg, lambda: avg_pool(x_bf16.to_mkldnn())
|
||||
)
|
||||
|
||||
def test_avg_pool2d_bf16(self):
|
||||
N = torch.randint(3, 10, (1,)).item()
|
||||
@ -586,13 +662,15 @@ class TestMkldnn(TestCase):
|
||||
kernel_size=3,
|
||||
stride=None,
|
||||
padding=1,
|
||||
count_include_pad=count_include_pad)
|
||||
count_include_pad=count_include_pad,
|
||||
)
|
||||
y2 = F.avg_pool2d(
|
||||
x.to_mkldnn(),
|
||||
kernel_size=3,
|
||||
stride=None,
|
||||
padding=1,
|
||||
count_include_pad=count_include_pad)
|
||||
count_include_pad=count_include_pad,
|
||||
)
|
||||
|
||||
self.assertEqual(y1, y2.to_dense())
|
||||
|
||||
@ -630,31 +708,34 @@ class TestMkldnn(TestCase):
|
||||
self.assertEqual(y, y_bf16, atol=1e-1, rtol=1e-3)
|
||||
else:
|
||||
msg = "mkldnn_adaptive_avg_pool2d: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: adaptive_avg_pool2d(x_bf16.to_mkldnn()))
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError, msg, lambda: adaptive_avg_pool2d(x_bf16.to_mkldnn())
|
||||
)
|
||||
|
||||
def _test_batch_norm_base(self, dim, channels, input):
|
||||
bn_module = {2 : torch.nn.BatchNorm2d, 3 : torch.nn.BatchNorm3d}
|
||||
bn_module = {2: torch.nn.BatchNorm2d, 3: torch.nn.BatchNorm3d}
|
||||
bn = bn_module[dim](channels).float().train(False)
|
||||
mkldnn_bn = mkldnn_utils.to_mkldnn(copy.deepcopy(bn))
|
||||
self.assertEqual(
|
||||
bn(input),
|
||||
mkldnn_bn(input.to_mkldnn()).to_dense())
|
||||
self.assertEqual(bn(input), mkldnn_bn(input.to_mkldnn()).to_dense())
|
||||
|
||||
self._test_serialization(mkldnn_bn, (input.to_mkldnn(),))
|
||||
self._test_tracing(mkldnn_bn, (input.to_mkldnn(),))
|
||||
|
||||
def _test_batch_norm_train_base(self, dim, channels, input):
|
||||
# TODO: support 3d batchnorm training.
|
||||
bn_module = {2 : torch.nn.BatchNorm2d}
|
||||
bn_module = {2: torch.nn.BatchNorm2d}
|
||||
# TODO: support none affine.
|
||||
options = itertools.product([True], [True, False])
|
||||
for affine, track_running_stats in options:
|
||||
bn = bn_module[dim](
|
||||
num_features=channels,
|
||||
affine=affine,
|
||||
track_running_stats=track_running_stats).float().train(True)
|
||||
bn = (
|
||||
bn_module[dim](
|
||||
num_features=channels,
|
||||
affine=affine,
|
||||
track_running_stats=track_running_stats,
|
||||
)
|
||||
.float()
|
||||
.train(True)
|
||||
)
|
||||
mkldnn_bn = copy.deepcopy(bn)
|
||||
x1 = input.clone().requires_grad_()
|
||||
x2 = input.clone().to_mkldnn().requires_grad_()
|
||||
@ -666,10 +747,14 @@ class TestMkldnn(TestCase):
|
||||
loss2.backward()
|
||||
self.assertEqual(y1, y2)
|
||||
self.assertEqual(x1.grad, x2.grad.to_dense())
|
||||
self.assertEqual(bn.weight.grad, mkldnn_bn.weight.grad, rtol=1e-3, atol=1e-3)
|
||||
self.assertEqual(
|
||||
bn.weight.grad, mkldnn_bn.weight.grad, rtol=1e-3, atol=1e-3
|
||||
)
|
||||
if track_running_stats:
|
||||
self.assertEqual(bn.running_mean, mkldnn_bn.running_mean)
|
||||
self.assertEqual(bn.running_var, mkldnn_bn.running_var, rtol=1e-5, atol=1e-5)
|
||||
self.assertEqual(
|
||||
bn.running_var, mkldnn_bn.running_var, rtol=1e-5, atol=1e-5
|
||||
)
|
||||
|
||||
def test_batch_norm_2d(self):
|
||||
N = torch.randint(3, 10, (1,)).item()
|
||||
@ -686,7 +771,7 @@ class TestMkldnn(TestCase):
|
||||
|
||||
@unittest.skipIf(IS_WINDOWS, "Limit support for bf16 path")
|
||||
def _test_batch_norm_bf16_base(self, dim, channels, input):
|
||||
bn_module = {2 : torch.nn.BatchNorm2d, 3 : torch.nn.BatchNorm3d}
|
||||
bn_module = {2: torch.nn.BatchNorm2d, 3: torch.nn.BatchNorm3d}
|
||||
x_bf16 = input.bfloat16()
|
||||
# TODO: support training
|
||||
for train in [False]:
|
||||
@ -698,9 +783,9 @@ class TestMkldnn(TestCase):
|
||||
self.assertEqual(y, y_bf16, atol=1e-1, rtol=1e-3)
|
||||
else:
|
||||
msg = "mkldnn_batch_norm: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: bn(x_bf16.to_mkldnn()))
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError, msg, lambda: bn(x_bf16.to_mkldnn())
|
||||
)
|
||||
|
||||
def test_batch_norm_2d_bf16(self):
|
||||
N = torch.randint(3, 10, (1,)).item()
|
||||
@ -725,13 +810,11 @@ class TestMkldnn(TestCase):
|
||||
my = y.to_mkldnn()
|
||||
|
||||
# add
|
||||
self.assertEqual(
|
||||
x + y,
|
||||
(mx + my).to_dense())
|
||||
self.assertEqual(x + y, (mx + my).to_dense())
|
||||
|
||||
self.assertEqual(
|
||||
torch.add(x, y, alpha=alpha),
|
||||
torch.add(mx, my, alpha=alpha).to_dense())
|
||||
torch.add(x, y, alpha=alpha), torch.add(mx, my, alpha=alpha).to_dense()
|
||||
)
|
||||
|
||||
# add_
|
||||
x += y
|
||||
@ -766,21 +849,13 @@ class TestMkldnn(TestCase):
|
||||
my = y.to_mkldnn()
|
||||
|
||||
# mul
|
||||
self.assertEqual(
|
||||
x * y,
|
||||
(mx * my).to_dense())
|
||||
self.assertEqual(x * y, (mx * my).to_dense())
|
||||
|
||||
self.assertEqual(
|
||||
x * value,
|
||||
(mx * value).to_dense())
|
||||
self.assertEqual(x * value, (mx * value).to_dense())
|
||||
|
||||
self.assertEqual(
|
||||
torch.mul(x, y),
|
||||
torch.mul(mx, my).to_dense())
|
||||
self.assertEqual(torch.mul(x, y), torch.mul(mx, my).to_dense())
|
||||
|
||||
self.assertEqual(
|
||||
torch.mul(x, value),
|
||||
torch.mul(mx, value).to_dense())
|
||||
self.assertEqual(torch.mul(x, value), torch.mul(mx, value).to_dense())
|
||||
|
||||
# mul_
|
||||
x *= y
|
||||
@ -838,9 +913,9 @@ class TestMkldnn(TestCase):
|
||||
|
||||
def test_view(self):
|
||||
x = torch.randn(3, 4, 5, dtype=torch.float32).to_mkldnn()
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
"Change to use reshape",
|
||||
lambda: x.view(x.size(0), -1))
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError, "Change to use reshape", lambda: x.view(x.size(0), -1)
|
||||
)
|
||||
|
||||
def test_reshape(self):
|
||||
x = torch.randn(3, 4, 5, dtype=torch.float32) * 10
|
||||
@ -941,9 +1016,7 @@ class TestMkldnn(TestCase):
|
||||
for bias in [True, False]:
|
||||
linear = torch.nn.Linear(in_features, out_features, bias=bias).float()
|
||||
mkldnn_linear = mkldnn_utils.to_mkldnn(copy.deepcopy(linear))
|
||||
self.assertEqual(
|
||||
linear(x),
|
||||
mkldnn_linear(x.to_mkldnn()).to_dense())
|
||||
self.assertEqual(linear(x), mkldnn_linear(x.to_mkldnn()).to_dense())
|
||||
|
||||
self._test_serialization(mkldnn_linear, (x.to_mkldnn(),))
|
||||
self._test_tracing(mkldnn_linear, (x.to_mkldnn(),))
|
||||
@ -976,24 +1049,24 @@ class TestMkldnn(TestCase):
|
||||
for bias in [True, False]:
|
||||
linear = torch.nn.Linear(in_features, out_features, bias=bias).float()
|
||||
mkldnn_linear = mkldnn_utils.to_mkldnn(copy.deepcopy(linear))
|
||||
mkldnn_linear_bf16 = mkldnn_utils.to_mkldnn(copy.deepcopy(linear), torch.bfloat16)
|
||||
mkldnn_linear_bf16 = mkldnn_utils.to_mkldnn(
|
||||
copy.deepcopy(linear), torch.bfloat16
|
||||
)
|
||||
if has_bf16_support():
|
||||
y = mkldnn_linear(x.to_mkldnn()).to_dense()
|
||||
y_bf16 = mkldnn_linear_bf16(x_bf16.to_mkldnn()).to_dense(torch.float32)
|
||||
self.assertEqual(y, y_bf16, atol=1e-1, rtol=1e-3)
|
||||
else:
|
||||
msg = "mkldnn_linear: bf16 path needs the cpu support avx512bw, avx512vl and avx512dq"
|
||||
self.assertRaisesRegex(RuntimeError,
|
||||
msg,
|
||||
lambda: mkldnn_linear_bf16(x_bf16.to_mkldnn()))
|
||||
self.assertRaisesRegex(
|
||||
RuntimeError, msg, lambda: mkldnn_linear_bf16(x_bf16.to_mkldnn())
|
||||
)
|
||||
|
||||
def test_softmax(self):
|
||||
x = torch.randn(3, 4, 5, dtype=torch.float32) * 10
|
||||
for dim in range(x.ndim):
|
||||
softmax = torch.nn.Softmax(dim=dim)
|
||||
self.assertEqual(
|
||||
softmax(x),
|
||||
softmax(x.to_mkldnn()).to_dense())
|
||||
self.assertEqual(softmax(x), softmax(x.to_mkldnn()).to_dense())
|
||||
|
||||
def test_sigmoid(self):
|
||||
x = torch.randn(4, 5, dtype=torch.float32) * 10
|
||||
@ -1023,22 +1096,18 @@ class TestMkldnn(TestCase):
|
||||
with TemporaryFileName() as fname:
|
||||
torch.jit.save(module, fname)
|
||||
loaded = torch.jit.load(fname)
|
||||
self.assertEqual(
|
||||
module(*inputs).to_dense(),
|
||||
loaded(*inputs).to_dense())
|
||||
self.assertEqual(module(*inputs).to_dense(), loaded(*inputs).to_dense())
|
||||
|
||||
def _test_tracing(self, module, inputs):
|
||||
traced = torch.jit.trace(module, inputs)
|
||||
self.assertEqual(
|
||||
module(*inputs).to_dense(),
|
||||
traced(*inputs).to_dense())
|
||||
self.assertEqual(module(*inputs).to_dense(), traced(*inputs).to_dense())
|
||||
|
||||
def test_set_data_tensorimpl_type(self):
|
||||
# Dense tensor has impl of type `TensorImpl`, while MKL-DNN tensor has impl
|
||||
# of type `OpaqueTensorImpl<IDeepTensorWrapperPtr>`.
|
||||
x = torch.randn((1, 2), dtype=torch.float, device=torch.device('cpu'))
|
||||
x = torch.randn((1, 2), dtype=torch.float, device=torch.device("cpu"))
|
||||
x_mkldnn = x.to_mkldnn()
|
||||
with self.assertRaisesRegex(RuntimeError, 'incompatible tensor type'):
|
||||
with self.assertRaisesRegex(RuntimeError, "incompatible tensor type"):
|
||||
x.data = x_mkldnn
|
||||
|
||||
def test_empty(self):
|
||||
@ -1064,7 +1133,7 @@ class TestMkldnn(TestCase):
|
||||
def test_legacy_new_failure(self):
|
||||
x = torch.randn(1, dtype=torch.float32)
|
||||
x_mkldnn = x.to_mkldnn()
|
||||
self.assertRaises(RuntimeError, lambda: x_mkldnn.new(device='cpu'))
|
||||
self.assertRaises(RuntimeError, lambda: x_mkldnn.new(device="cpu"))
|
||||
self.assertRaises(RuntimeError, lambda: x_mkldnn.new(x.storage()))
|
||||
self.assertRaises(RuntimeError, lambda: x_mkldnn.new(x))
|
||||
self.assertRaises(RuntimeError, lambda: x_mkldnn.new(torch.Size([2, 3])))
|
||||
@ -1104,5 +1173,5 @@ class TestMkldnn(TestCase):
|
||||
self._test_imagenet_model(model)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
run_tests()
|
||||
|
||||
Reference in New Issue
Block a user