mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
bf7e290854
This PR is part of a series attempting to re-submit https://github.com/pytorch/pytorch/pull/134592 as smaller PRs. In jit tests: - Add and use a common raise_on_run_directly method for when a user runs a test file directly which should not be run this way. Print the file which the user should have run. - Raise a RuntimeError on tests which have been disabled (not run) Pull Request resolved: https://github.com/pytorch/pytorch/pull/154725 Approved by: https://github.com/clee2000
818 lines
29 KiB
Python
818 lines
29 KiB
Python
# Owner(s): ["oncall: jit"]
|
|
|
|
import operator
|
|
import unittest
|
|
from textwrap import dedent
|
|
from typing import Any, List
|
|
|
|
import torch
|
|
from torch import nn, Tensor
|
|
from torch.testing import FileCheck
|
|
from torch.testing._internal.common_methods_invocations import sample_inputs_cat_concat
|
|
from torch.testing._internal.common_utils import make_tensor, raise_on_run_directly
|
|
from torch.testing._internal.jit_utils import execWrapper, JitTestCase
|
|
|
|
|
|
# XXX: still in prototype
|
|
class TestSymbolicShapeAnalysis(JitTestCase):
|
|
def setUp(self):
|
|
super(JitTestCase, self).setUp()
|
|
self.prev_symbolic_shapes_test_enabled = (
|
|
torch._C._jit_symbolic_shapes_test_mode_enabled()
|
|
)
|
|
torch._C._jit_set_symbolic_shapes_test_mode(True)
|
|
|
|
def tearDown(self):
|
|
torch._C._jit_set_symbolic_shapes_test_mode(
|
|
self.prev_symbolic_shapes_test_enabled
|
|
)
|
|
|
|
def test_shape_analysis(self):
|
|
@torch.jit.script
|
|
def foo(x, y):
|
|
return x * y
|
|
|
|
inputs = list(foo.graph.inputs())
|
|
|
|
def prop_shapes_on_graph(inp0, inp1):
|
|
inputs[0].setType(inputs[0].type().with_sizes(inp0))
|
|
inputs[1].setType(inputs[1].type().with_sizes(inp1))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(foo.graph)
|
|
|
|
prop_shapes_on_graph([1, 6, 5], [1, 7, 1, 5])
|
|
FileCheck().check("1, 7, 6, 5").run(foo.graph)
|
|
|
|
# None implicitly creates a new symbolic symbol
|
|
prop_shapes_on_graph([None, None], [None, None, None])
|
|
output_shape = foo.graph.findNode("aten::mul").output().type().symbolic_sizes()
|
|
inp0_shape = inputs[0].type().symbolic_sizes()
|
|
inp1_shape = inputs[1].type().symbolic_sizes()
|
|
|
|
# output shape dim 0 should be taken from the second inp dim0
|
|
# other two dims we cannot infer and are given a new symbolic shape
|
|
self.assertEqual(output_shape[0], inp1_shape[0])
|
|
self.assertFalse(output_shape[1] in inp0_shape + inp1_shape)
|
|
self.assertFalse(output_shape[2] in inp0_shape + inp1_shape)
|
|
|
|
# XXX: symbolic shapes are represented with an increasing counter of unique
|
|
# values, use `_new_symbolic_shape_symbol` api instead of specifying negative
|
|
# dimensions directly so there is no chance of collision between manual number
|
|
# and current counter value.
|
|
sym1 = torch._C._new_symbolic_shape_symbol()
|
|
sym2 = torch._C._new_symbolic_shape_symbol()
|
|
sym3 = torch._C._new_symbolic_shape_symbol()
|
|
prop_shapes_on_graph([sym1, 1, sym3], [1, sym2, sym3])
|
|
output_shape = foo.graph.findNode("aten::mul").output().type().symbolic_sizes()
|
|
self.assertEqual(output_shape[0], sym1)
|
|
self.assertEqual(output_shape[1], sym2)
|
|
self.assertEqual(output_shape[2], sym3)
|
|
|
|
def test_shared_shape_graph(self):
|
|
@torch.jit.script
|
|
def foo(x, y):
|
|
return x * y, x / y
|
|
|
|
mul_node = foo.graph.findNode("aten::mul")
|
|
div_node = foo.graph.findNode("aten::div")
|
|
|
|
mul_graph = torch._C._jit_shape_compute_graph_for_node(mul_node)
|
|
div_graph = torch._C._jit_shape_compute_graph_for_node(div_node)
|
|
self.assertIsNotNone(mul_graph)
|
|
self.assertIs(mul_graph, div_graph)
|
|
|
|
def test_write(self):
|
|
@torch.jit.script
|
|
def foo(a, b):
|
|
return a * b
|
|
|
|
# broadcast appends cant be removed, so we bail on propagation
|
|
torch._C._jit_pass_propagate_shapes_on_graph(foo.graph)
|
|
FileCheck().check("Tensor = aten::mul").run(foo.graph)
|
|
|
|
@torch.jit.script
|
|
def foo(y):
|
|
x = [1, 2, 3, 4]
|
|
x[0] = 5
|
|
return y.view(x)
|
|
|
|
torch._C._jit_pass_propagate_shapes_on_graph(foo.graph)
|
|
FileCheck().check("Tensor = aten::view").run(foo.graph)
|
|
|
|
def test_if_propagation(self):
|
|
@torch.jit.script
|
|
def foo(i: int, z):
|
|
x = torch.ones([2, 3, 4, 5])
|
|
y = z.view([z.size(i), 3, 2, z.size(i)])
|
|
if i == 4:
|
|
return x
|
|
else:
|
|
return y
|
|
|
|
torch._C._jit_pass_constant_propagation(foo.graph)
|
|
torch._C._jit_pass_propagate_shapes_on_graph(foo.graph)
|
|
view = foo.graph.findNode("aten::view")
|
|
|
|
def neg_to_one(li):
|
|
return [elem if elem >= 0 else -1 for elem in li]
|
|
|
|
self.assertEqual(
|
|
neg_to_one(view.output().type().symbolic_sizes()), [-1, 3, 2, -1]
|
|
)
|
|
if_out = next(foo.graph.findNode("prim::If").outputs())
|
|
self.assertEqual(neg_to_one(if_out.type().symbolic_sizes()), [-1, 3, -1, -1])
|
|
|
|
def test_unary_shape_functions(self):
|
|
unary_ops = [
|
|
torch.nn.functional.hardtanh,
|
|
]
|
|
for fn in unary_ops:
|
|
t = torch.jit.trace(fn, (torch.rand([4, 4])))
|
|
ten_input = next(t.graph.inputs())
|
|
ten_input.setType(ten_input.type().with_sizes([2, 2]))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(t.graph)
|
|
self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [2, 2])
|
|
|
|
def test_unary_shape_fns_inplace(self):
|
|
def mul_inplace(x: torch.Tensor):
|
|
y = x.mul_(2)
|
|
return y
|
|
|
|
unary_ops = [mul_inplace]
|
|
for fn in unary_ops:
|
|
# t = torch.jit.trace(fn, torch.rand([4, 4])) # For some reason tracing is erroring out.
|
|
t = torch.jit.script(fn)
|
|
ten_input = next(t.graph.inputs())
|
|
ten_input.setType(ten_input.type().with_sizes([2, 2]))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(t.graph)
|
|
self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [2, 2])
|
|
|
|
def test_binary_shape_functions(self):
|
|
binary_ops = [
|
|
operator.__mul__,
|
|
operator.__truediv__,
|
|
operator.__gt__,
|
|
operator.__add__,
|
|
]
|
|
|
|
for fn in binary_ops:
|
|
size_1 = [1, 4, 8]
|
|
size_2 = [4, 1, 8]
|
|
t = torch.jit.trace(fn, (torch.rand([4]), torch.rand([4])))
|
|
inputs = list(t.graph.inputs())
|
|
inputs[0].setType(inputs[0].type().with_sizes(size_1))
|
|
inputs[1].setType(inputs[1].type().with_sizes(size_2))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(t.graph)
|
|
self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [4, 4, 8])
|
|
|
|
def test_binary_shape_fns_inplace(self):
|
|
def div_inplace_tensor(x: torch.Tensor, y: torch.Tensor):
|
|
z = x.div_(y)
|
|
return z
|
|
|
|
def add_inplace_tensor(x: torch.Tensor, y: torch.Tensor):
|
|
z = x.add_(y)
|
|
return z
|
|
|
|
binary_ops = [
|
|
div_inplace_tensor,
|
|
add_inplace_tensor,
|
|
]
|
|
|
|
for fn in binary_ops:
|
|
size_1 = [4, 4, 8] # x (can't broadcast because it's an inplace op)
|
|
t = torch.jit.script(fn)
|
|
inputs = list(t.graph.inputs())
|
|
inputs[0].setType(inputs[0].type().with_sizes(size_1))
|
|
# Intentionally not populate the type of inputs[1]
|
|
torch._C._jit_pass_propagate_shapes_on_graph(t.graph)
|
|
self.assertEqual(next(t.graph.outputs()).type().symbolic_sizes(), [4, 4, 8])
|
|
|
|
def test_size_and_sizes(self):
|
|
@torch.jit.script
|
|
def foo(x, y):
|
|
return x.view(y.size(0), 8, y.size(-1))
|
|
|
|
@torch.jit.script
|
|
def foo2(x, y):
|
|
return x.view(y.size())
|
|
|
|
for graph in [foo.graph, foo2.graph]:
|
|
inputs = list(graph.inputs())
|
|
sym1 = torch._C._new_symbolic_shape_symbol()
|
|
|
|
inputs[1].setType(inputs[1].type().with_sizes([5, 8, sym1]))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(graph)
|
|
self.assertEqual(
|
|
next(graph.outputs()).type().symbolic_sizes(), [5, 8, sym1]
|
|
)
|
|
|
|
def test_adaptive_avg_pool2d(self):
|
|
inps = [
|
|
[(1, 64, 8, 9), (5, 7)],
|
|
[(1, 64, 10, 9), (7)],
|
|
[(1, 64, 10, 9), (5, None)],
|
|
[(1, 8, 4, 3), (None, None)],
|
|
[(1, 8, 4, 3), (None, 5)],
|
|
]
|
|
|
|
for inp in inps:
|
|
t = torch.randn(*inp[0])
|
|
out_size = torch.nn.functional.adaptive_avg_pool2d(t, inp[1]).size()
|
|
|
|
def foo(x):
|
|
return torch.nn.functional.adaptive_avg_pool2d(x, inp[1])
|
|
|
|
fn = torch.jit.trace(foo, (t,))
|
|
torch._C._jit_erase_non_input_shape_information(fn.graph)
|
|
torch._C._jit_pass_peephole(fn.graph)
|
|
torch._C._jit_pass_constant_propagation(fn.graph)
|
|
self.checkShapeAnalysis(out_size, fn.graph, assert_propagation=True)
|
|
|
|
def test_conv_deconv(self):
|
|
for (
|
|
inp_shape,
|
|
weight_shape,
|
|
bias,
|
|
stride,
|
|
padding,
|
|
output_padding,
|
|
dilation,
|
|
groups,
|
|
mod,
|
|
) in [
|
|
([32, 6, 10], [16, 3, 3], None, 2, 2, 1, 1, 2, torch.nn.functional.conv1d),
|
|
(
|
|
[32, 16, 10],
|
|
[16, 3, 3],
|
|
None,
|
|
2,
|
|
2,
|
|
1,
|
|
1,
|
|
2,
|
|
torch.nn.functional.conv_transpose1d,
|
|
),
|
|
(
|
|
[1, 32, 5, 10],
|
|
[30, 16, 3, 3],
|
|
None,
|
|
[2, 2],
|
|
[0, 0],
|
|
0,
|
|
1,
|
|
2,
|
|
torch.nn.functional.conv2d,
|
|
),
|
|
(
|
|
[1, 30, 5, 10],
|
|
[30, 16, 3, 3],
|
|
None,
|
|
[2, 2],
|
|
[0, 0],
|
|
0,
|
|
1,
|
|
2,
|
|
torch.nn.functional.conv_transpose2d,
|
|
),
|
|
(
|
|
[3, 14, 10, 66, 55],
|
|
[2, 7, 7, 4, 4],
|
|
None,
|
|
1,
|
|
1,
|
|
2,
|
|
1,
|
|
2,
|
|
torch.nn.functional.conv3d,
|
|
),
|
|
(
|
|
[3, 2, 10, 66, 55],
|
|
[2, 7, 7, 4, 4],
|
|
None,
|
|
1,
|
|
1,
|
|
0,
|
|
1,
|
|
2,
|
|
torch.nn.functional.conv_transpose3d,
|
|
),
|
|
]:
|
|
inp = torch.rand(inp_shape)
|
|
weight = torch.rand(weight_shape)
|
|
if mod in [
|
|
torch.nn.functional.conv1d,
|
|
torch.nn.functional.conv2d,
|
|
torch.nn.functional.conv3d,
|
|
]:
|
|
res = mod(inp, weight, bias, stride, padding, dilation, groups).size()
|
|
else:
|
|
res = mod(
|
|
inp, weight, bias, stride, padding, output_padding, dilation, groups
|
|
).size()
|
|
|
|
def foo(inp, weight):
|
|
if mod in [
|
|
torch.nn.functional.conv1d,
|
|
torch.nn.functional.conv2d,
|
|
torch.nn.functional.conv3d,
|
|
]:
|
|
return mod(inp, weight, bias, stride, padding, dilation, groups)
|
|
else:
|
|
return mod(
|
|
inp,
|
|
weight,
|
|
bias,
|
|
stride,
|
|
padding,
|
|
output_padding,
|
|
dilation,
|
|
groups,
|
|
)
|
|
|
|
fn = torch.jit.trace(foo, (inp, weight))
|
|
torch._C._jit_erase_non_input_shape_information(fn.graph)
|
|
torch._C._jit_pass_peephole(fn.graph)
|
|
torch._C._jit_pass_constant_propagation(fn.graph)
|
|
self.checkShapeAnalysis(res, fn.graph, assert_propagation=True)
|
|
|
|
def test_arange_shape(self):
|
|
# no opinfo for tensor constructors
|
|
inps = [
|
|
(10,),
|
|
(10, 10),
|
|
(0, 10),
|
|
(0, 1000),
|
|
(1, -1, -1),
|
|
(1, 0, -1),
|
|
(1, 2, 1),
|
|
(0.6, 0.89, 0.1),
|
|
(1, 10, 0.3),
|
|
(1, 10, 4),
|
|
(0.6, 0.7, 0.8),
|
|
(1, 10, 0.3),
|
|
# (True,), TODO: https://github.com/pytorch/pytorch/issues/63405
|
|
# (False,), TODO: https://github.com/pytorch/pytorch/issues/63405
|
|
(0, 5),
|
|
(0, 5, 2),
|
|
(0, 5 + 1e-6),
|
|
(0, 5 - 1e-6),
|
|
(10, -1 + 1e-6, -1),
|
|
(10, -1, -1),
|
|
(10, -1 - 1e-6, -1),
|
|
]
|
|
|
|
for inp in inps:
|
|
funcs_template = dedent(
|
|
"""
|
|
def func():
|
|
return torch.arange({args})
|
|
"""
|
|
)
|
|
|
|
inp_s = str(inp)[1:-1] # remove tuple parens
|
|
funcs_str = funcs_template.format(args=inp_s)
|
|
scope = {}
|
|
execWrapper(funcs_str, globals(), scope)
|
|
cu = torch.jit.CompilationUnit(funcs_str)
|
|
self.checkShapeAnalysis(
|
|
list(cu.func().size()),
|
|
cu.func.graph,
|
|
assert_propagation=True,
|
|
constant_prop=False,
|
|
)
|
|
|
|
def test_shape_embedding_bag(self):
|
|
# TODO: merge into opinfos, having difficulties there
|
|
with torch.no_grad():
|
|
|
|
def make_arg(shape, low=None, high=None):
|
|
return make_tensor(
|
|
shape,
|
|
device="cpu",
|
|
dtype=torch.int64,
|
|
low=low,
|
|
high=high,
|
|
requires_grad=False,
|
|
)
|
|
|
|
nn_inps = (
|
|
(
|
|
make_arg((40,), 0, 9),
|
|
torch.nn.Embedding(20, embedding_dim=64, max_norm=1.0),
|
|
),
|
|
(make_arg((2, 4), 0, 9), torch.nn.Embedding(10, 20, sparse=True)),
|
|
(make_arg((0,)), torch.nn.Embedding(0, 0, sparse=True)),
|
|
(make_arg((2, 4), 0, 9), torch.nn.Embedding(10, 0, sparse=True)),
|
|
(make_arg((4,), 0, 21), torch.nn.Embedding(22, 5, max_norm=1.0)),
|
|
(
|
|
make_arg((2,), 0, 1),
|
|
torch.nn.Embedding.from_pretrained(
|
|
torch.arange(6.0).view(2, 3),
|
|
max_norm=2.0,
|
|
norm_type=0.5,
|
|
scale_grad_by_freq=False,
|
|
sparse=True,
|
|
),
|
|
),
|
|
)
|
|
|
|
for inp, module in nn_inps:
|
|
kwargs = {
|
|
"weight": module.weight.detach(),
|
|
"padding_idx": module.padding_idx,
|
|
"max_norm": module.max_norm,
|
|
"norm_type": module.norm_type,
|
|
"scale_grad_by_freq": module.scale_grad_by_freq,
|
|
"sparse": module.sparse,
|
|
}
|
|
|
|
out_size = torch.nn.functional.embedding(inp, **kwargs).size()
|
|
|
|
def foo(x):
|
|
return torch.nn.functional.embedding(inp, **kwargs)
|
|
|
|
fn = torch.jit.trace(foo, (inp.detach(),), check_trace=False)
|
|
|
|
self.checkShapeAnalysis(
|
|
out_size, fn.graph, assert_propagation=True, constant_prop=False
|
|
)
|
|
|
|
def test_shape_concat(self):
|
|
# TODO: unify with opinfo tests, traces of lists dont preserve sizes in IR
|
|
sample_inputs = sample_inputs_cat_concat(None, "cpu", torch.float, False)
|
|
|
|
class CatMod(nn.Module):
|
|
__constants__ = ["dim"]
|
|
|
|
def __init__(self, dim=0):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x, y):
|
|
return torch.cat([x, y], dim=self.dim)
|
|
|
|
for inp in sample_inputs:
|
|
mod = torch.jit.script(CatMod(**inp.kwargs).eval())
|
|
|
|
args = inp.input
|
|
|
|
# This test is hard-coded only to work with two sample inputs
|
|
# but the OpInfo may have more/less
|
|
if len(args) != 2:
|
|
continue
|
|
|
|
out_size = mod(*args).size()
|
|
inps = list(mod.graph.inputs())
|
|
inps[1].setType(inps[1].type().with_sizes(args[0].size()))
|
|
inps[2].setType(inps[2].type().with_sizes(args[1].size()))
|
|
self.checkShapeAnalysis(out_size, mod.graph, assert_propagation=True)
|
|
|
|
def assert_shape_equal_scripted(self, script_fn, given_ins):
|
|
expected_res = script_fn(*given_ins)
|
|
g = script_fn.graph
|
|
graph_ins = list(g.inputs())
|
|
self.assertEqual(len(given_ins), len(graph_ins))
|
|
for inp, graph_in in zip(given_ins, graph_ins):
|
|
graph_in.setType(graph_in.type().with_sizes(inp.size()))
|
|
|
|
out_sizes = [out.size() for out in expected_res]
|
|
self.checkShapeAnalysis(out_sizes, g, assert_propagation=True)
|
|
|
|
def test_convolution_backward(self):
|
|
# No opinfos for ops that are not part of the Python API
|
|
# Also, as the return shapes are the input, weight, and bias shape, there is no point
|
|
# in a really complicated test
|
|
|
|
input = torch.randn(
|
|
(16, 16, 8, 8), dtype=torch.float32, device="cpu", requires_grad=True
|
|
)
|
|
weight = torch.randn(
|
|
(8, 4, 3, 3), dtype=torch.float32, device="cpu", requires_grad=True
|
|
)
|
|
out_grad = torch.randn((16, 8, 8, 8), dtype=torch.float32, device="cpu")
|
|
|
|
@torch.jit.script
|
|
def conv_bwd(input, weight, grad):
|
|
bias_sizes = [
|
|
8,
|
|
]
|
|
args = ([1, 1], [1, 1], [1, 1], False, [0, 0], 4, [True, True, True])
|
|
return torch.ops.aten.convolution_backward(
|
|
grad, input, weight, bias_sizes, *args
|
|
)
|
|
|
|
self.assert_shape_equal_scripted(conv_bwd, (input, weight, out_grad))
|
|
|
|
@torch.jit.script
|
|
def conv_bwd_2(input, weight, grad):
|
|
bias_sizes = None
|
|
args = ([1, 1], [1, 1], [1, 1], False, [0, 0], 4, [True, True, True])
|
|
return torch.ops.aten.convolution_backward(
|
|
grad, input, weight, bias_sizes, *args
|
|
)
|
|
|
|
self.assert_shape_equal_scripted(conv_bwd_2, (input, weight, out_grad))
|
|
|
|
def test_returning_input_symbolic_shapes(self):
|
|
mm = torch.jit.freeze(torch.jit.script(nn.Conv2d(16, 33, 3, stride=2).eval()))
|
|
inps = list(mm.graph.inputs())
|
|
inps[1].setType(inps[1].type().with_sizes([None, None, None, None]))
|
|
shape_compute_graph = (
|
|
torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mm.graph)
|
|
)
|
|
g = shape_compute_graph.partial_eval_shape_graph()
|
|
# to make into a jit function cant have multiple outputs
|
|
g.makeMultiOutputIntoTuple()
|
|
func = torch._C._create_function_from_graph("partial_eval_graph", g)
|
|
out = func([20, 16, 5, 10])
|
|
# first four outputs should be unknown symbolic shapes from input
|
|
self.assertEqual(out[0:4], [20, 16, 5, 10])
|
|
# last two are two new symbolic dims - height and width
|
|
self.assertEqual(out[4:], list(mm(torch.rand([20, 16, 5, 10])).size()[2:]))
|
|
|
|
def test_partial_eval_graph_conv(self):
|
|
mm = torch.jit.freeze(torch.jit.script(nn.Conv2d(16, 33, 3, stride=2).eval()))
|
|
shape_compute_graph = (
|
|
torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mm.graph)
|
|
)
|
|
output_sizes = (
|
|
mm.graph.findNode("aten::conv2d").output().type().symbolic_sizes()
|
|
)
|
|
# calculating 0, 2 and 3 index
|
|
for i in [0, 2, 3]:
|
|
self.assertTrue(output_sizes[i] < 0)
|
|
self.assertTrue(output_sizes[1] >= 0)
|
|
g = shape_compute_graph.partial_eval_shape_graph()
|
|
# to make into a jit function cant have multiple outputs
|
|
g.makeMultiOutputIntoTuple()
|
|
func = torch._C._create_function_from_graph("partial_eval_graph", g)
|
|
inp = torch.randn(20, 16, 5, 10)
|
|
output = func([20, 16, 5, 10])
|
|
output_eager = list(mm(inp).size())
|
|
for o, oe in zip(output, output_eager[0:1] + output_eager[2:]):
|
|
self.assertEqual(o, oe)
|
|
|
|
def checkSymShapeCompute(
|
|
self, shape_compute_graph, nodes, node_output_sizes, shape_inputs
|
|
):
|
|
g = shape_compute_graph.partial_eval_shape_graph()
|
|
self.assertTrue(len(list(g.inputs())) == len(shape_inputs))
|
|
output_sym_map = shape_compute_graph.graph_output_to_symbolic_shape_dim()
|
|
# map from sym shape -> index
|
|
sym_shape_to_index = {}
|
|
for index, output in enumerate(g.outputs()):
|
|
sym_shape_to_index[output_sym_map[output]] = index
|
|
|
|
g.makeMultiOutputIntoTuple()
|
|
func = torch._C._create_function_from_graph("partial_eval_graph", g)
|
|
sym_outputs = func(*shape_inputs)
|
|
|
|
for node, output_shape in zip(nodes, node_output_sizes):
|
|
output_type_sizes = node.output().type().symbolic_sizes()
|
|
for i, sym_shape in enumerate(output_type_sizes):
|
|
if sym_shape >= 0:
|
|
self.assertEqual(sym_shape, output_shape[i])
|
|
else:
|
|
sym_shape_index = sym_shape_to_index[sym_shape]
|
|
self.assertEqual(sym_outputs[sym_shape_index], output_shape[i])
|
|
|
|
def test_partial_eval_stitching(self):
|
|
conv1 = torch.nn.Conv2d(
|
|
3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
|
|
)
|
|
max_pool = torch.nn.MaxPool2d(
|
|
kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False
|
|
)
|
|
conv2 = nn.Conv2d(
|
|
64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False
|
|
)
|
|
|
|
mod = torch.jit.freeze(
|
|
torch.jit.script(nn.Sequential(conv1, max_pool, conv2).eval())
|
|
)
|
|
|
|
conv1_output = conv1(torch.rand(1, 3, 224, 224))
|
|
max_pool_output = max_pool(conv1_output)
|
|
conv2_output = conv2(max_pool_output)
|
|
|
|
shape_compute_graph = (
|
|
torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mod.graph)
|
|
)
|
|
nodes = [mod.graph.findNode("aten::max_pool2d")] + list(
|
|
mod.graph.findAllNodes("aten::conv2d")
|
|
)
|
|
output_shapes = [
|
|
max_pool_output.size(),
|
|
conv1_output.size(),
|
|
conv2_output.size(),
|
|
]
|
|
self.checkSymShapeCompute(
|
|
shape_compute_graph, nodes, output_shapes, ([1, 3, 224, 224],)
|
|
)
|
|
|
|
def test_refinement_through_graph_stitching(self):
|
|
class TwoConvs(torch.nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.conv1 = torch.nn.Conv2d(
|
|
3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
|
|
)
|
|
self.conv2 = torch.nn.Conv2d(
|
|
3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
|
|
)
|
|
|
|
def forward(self, x):
|
|
a = self.conv1(x)
|
|
b = self.conv2(x)
|
|
return a + b
|
|
|
|
mod = torch.jit.freeze(torch.jit.script(TwoConvs()).eval())
|
|
inp_tensor = list(mod.graph.inputs())[1]
|
|
inp_tensor.setType(inp_tensor.type().with_sizes([None, None, None, None]))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(mod.graph)
|
|
outs = list(next(mod.graph.outputs()).node().inputs())
|
|
out1 = outs[0].type().symbolic_sizes()
|
|
out2 = outs[1].type().symbolic_sizes()
|
|
self.assertTrue(out1[2] != out2[2])
|
|
self.assertTrue(out1[3] != out2[3])
|
|
# by joining partial eval graphs of both convs we are able to recognize the output shapes
|
|
# are equivalent
|
|
torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mod.graph)
|
|
out1 = outs[0].type().symbolic_sizes()
|
|
out2 = outs[1].type().symbolic_sizes()
|
|
self.assertEqual(out1, out2)
|
|
|
|
def test_stitching_multi_output(self):
|
|
max_pool = torch.nn.MaxPool2d(
|
|
kernel_size=3,
|
|
stride=2,
|
|
padding=1,
|
|
dilation=1,
|
|
ceil_mode=False,
|
|
return_indices=True,
|
|
)
|
|
tensor = torch.rand(1, 3, 224, 224)
|
|
mod = torch.jit.trace(max_pool, (tensor,))
|
|
mod = torch.jit.freeze(mod.eval())
|
|
inp = list(mod.graph.inputs())[1]
|
|
inp.setType(inp.type().with_sizes([None, None, None, None]))
|
|
output_tensor = list(mod(tensor)[0].size())
|
|
self.run_pass("lower_all_tuples", mod.graph)
|
|
shape_compute_graph = (
|
|
torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(mod.graph)
|
|
)
|
|
max_pool_node = mod.graph.findNode("aten::max_pool2d_with_indices")
|
|
outs = list(max_pool_node.outputs())
|
|
self.assertEqual(
|
|
outs[0].type().symbolic_sizes(), outs[1].type().symbolic_sizes()
|
|
)
|
|
g = shape_compute_graph.partial_eval_shape_graph()
|
|
# to make into a jit function cant have multiple outputs
|
|
g.makeMultiOutputIntoTuple()
|
|
func = torch._C._create_function_from_graph("partial_eval_graph", g)
|
|
mapping = shape_compute_graph.graph_output_to_symbolic_shape_dim() # noqa: F841
|
|
output_shape = func(tensor.size())
|
|
# the first 4 dims are input sym dimensions, then the ,
|
|
self.assertEqual(list(output_shape[0:4]), list(tensor.size()))
|
|
self.assertEqual(list(output_shape[4:]), output_tensor[2:])
|
|
|
|
def test_sym_ir_parsing(self):
|
|
graph_str1 = """graph(%x.1 : Float(SS(-2), SS(-3))):
|
|
%3 : int = prim::Constant[value=1]()
|
|
%4 : Tensor = aten::add(%x.1, %x.1, %3)
|
|
return (%4)"""
|
|
g = torch._C.parse_ir(graph_str1)
|
|
inp = next(g.inputs())
|
|
out = inp.type().symbolic_sizes()
|
|
self.assertEqual(out, [-2, -3])
|
|
|
|
def test_stitching_concat(self):
|
|
@torch.jit.script
|
|
def foo1(a, b, x, y):
|
|
return (a / b) + torch.cat([x, y])
|
|
|
|
@torch.jit.script
|
|
def foo2(a, b, x, y):
|
|
return (a / b) + torch.cat([x, y], dim=-2)
|
|
|
|
for foo in [foo1, foo2]:
|
|
g = foo.graph
|
|
for inp in foo.graph.inputs():
|
|
inp.setType(inp.type().with_sizes([None, None]))
|
|
|
|
shape_compute_graph = (
|
|
torch._C._jit_pass_propagate_shapes_on_graph_and_build_compute(
|
|
foo.graph
|
|
)
|
|
)
|
|
nodes = (
|
|
[g.findNode("aten::div")]
|
|
+ [g.findNode("aten::add")]
|
|
+ [g.findNode("aten::cat")]
|
|
)
|
|
|
|
inps = [1, 10], [20, 10], [15, 1], [5, 1]
|
|
output_shapes = [[20, 10], [20, 10], [20, 1]]
|
|
|
|
self.checkSymShapeCompute(shape_compute_graph, nodes, output_shapes, inps)
|
|
|
|
@unittest.skipIf(
|
|
not hasattr(torch.jit, "_shapes"), "shape functions not loaded in python"
|
|
)
|
|
def test_shape_function_includes(self):
|
|
inp_shape = [1, 16, 5, 10]
|
|
weight_shape = [33, 16, 3, 3]
|
|
bias = None
|
|
stride = [2, 2]
|
|
padding = [0, 0]
|
|
dilation = [1, 1]
|
|
groups = 1
|
|
res = torch.jit._shapes.conv2d(
|
|
inp_shape, weight_shape, bias, stride, padding, dilation, groups
|
|
)
|
|
self.assertEqual(res, [1, 33, 2, 4])
|
|
|
|
m1_shape = [10, 20]
|
|
m2_shape = [20, 10]
|
|
res = torch.jit._shapes.matmul(m1_shape, m2_shape)
|
|
self.assertEqual(res, [10, 10])
|
|
|
|
def test_register_function_error_checking(self):
|
|
# this will error before registering on global map, so
|
|
# no issue in overwriting schema mappings
|
|
@torch.jit.script
|
|
def foo(x, y):
|
|
return x + y
|
|
|
|
node = foo.graph.findNode("aten::add")
|
|
|
|
@torch.jit.script
|
|
def wrong_input_types(x, y):
|
|
x: List[int] = []
|
|
return x
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "Expected supertype of int"):
|
|
torch._C._jit_register_shape_compute_graph_for_node(
|
|
node, wrong_input_types.graph
|
|
)
|
|
|
|
@torch.jit.script
|
|
def wrong_output_types(x: List[int], y: List[int]):
|
|
x: List[Tensor] = []
|
|
return x
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "but got graph_type"):
|
|
torch._C._jit_register_shape_compute_graph_for_node(
|
|
node, wrong_output_types.graph
|
|
)
|
|
|
|
@torch.jit.script
|
|
def too_many_inputs(x: List[int], y: List[int], z: Any, z2: Any):
|
|
x: List[int] = []
|
|
return x
|
|
|
|
with self.assertRaises(RuntimeError) as error:
|
|
torch._C._jit_register_shape_compute_graph_for_node(
|
|
node, too_many_inputs.graph
|
|
)
|
|
|
|
self.assertTrue("fewer arguments than schema" in str(error.exception))
|
|
|
|
def test_cross_entropy_loss(self):
|
|
@torch.jit.script
|
|
def foo(x, y):
|
|
return torch.ops.aten.cross_entropy_loss(x, y, reduction=0)
|
|
|
|
inputs = list(foo.graph.inputs())
|
|
inputs[0].setType(inputs[0].type().with_sizes([8, 2]))
|
|
inputs[1].setType(
|
|
inputs[1]
|
|
.type()
|
|
.with_sizes(
|
|
[
|
|
8,
|
|
]
|
|
)
|
|
)
|
|
torch._C._jit_pass_propagate_shapes_on_graph(foo.graph)
|
|
self.assertEqual(
|
|
next(foo.graph.outputs()).type().sizes(),
|
|
[
|
|
8,
|
|
],
|
|
)
|
|
|
|
def test_squeeze_dims(self):
|
|
@torch.jit.script
|
|
def foo(x):
|
|
return torch.ops.aten.squeeze(x, dim=0)
|
|
|
|
input = next(foo.graph.inputs())
|
|
input.setType(input.type().with_sizes([1, 5, 8]))
|
|
torch._C._jit_pass_propagate_shapes_on_graph(foo.graph)
|
|
self.assertEqual(next(foo.graph.outputs()).type().symbolic_sizes(), [5, 8])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise_on_run_directly("test/test_jit.py")
|