def test_find_and_run_method(self):
class MyTestModule(torch.nn.Module):
def forward(self, arg):
return arg
input = (torch.tensor([1]), )
script_module = torch.jit.script(MyTestModule())
script_module_result = script_module(*input)
buffer = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
has_bundled_inputs = mobile_module.find_method(
"get_all_bundled_inputs")
self.assertFalse(has_bundled_inputs)
torch.utils.bundled_inputs.augment_model_with_bundled_inputs(
script_module, [input], [])
buffer = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
has_bundled_inputs = mobile_module.find_method(
"get_all_bundled_inputs")
self.assertTrue(has_bundled_inputs)
bundled_inputs = mobile_module.run_method("get_all_bundled_inputs")
mobile_module_result = mobile_module.forward(*bundled_inputs[0])
torch.testing.assert_close(script_module_result, mobile_module_result)
def test_versioned_div_scalar_inplace(self, sample_input):
def historic_div_scalar_float_inplace(self, other: float):
return self.true_divide_(other)
def historic_div_scalar_int_inplace(self, other: int):
if self.is_floating_point():
return self.true_divide_(other)
return self.divide_(other, rounding_mode='trunc')
class MyModuleFloat(torch.nn.Module):
def __init__(self):
super(MyModuleFloat, self).__init__()
def forward(self, a, b: float):
a /= b
return a
class MyModuleInt(torch.nn.Module):
def __init__(self):
super(MyModuleInt, self).__init__()
def forward(self, a, b: int):
a /= b
return a
try:
v3_mobile_module_float = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_scalar_inplace_float_v2.ptl"
)
v3_mobile_module_int = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_scalar_inplace_int_v2.ptl"
)
except Exception as e:
self.skipTest("Failed to load fixture!")
current_mobile_module_float = self._save_load_module(MyModuleFloat)
current_mobile_module_int = self._save_load_module(MyModuleInt)
for val_a, val_b in product(sample_input, sample_input):
a = torch.tensor((val_a, ))
b = val_b
def _helper(m, fn):
m_result = self._try_fn(m, a, b)
fn_result = self._try_fn(fn, a, b)
if isinstance(m_result, Exception):
self.assertTrue(fn_result, Exception)
else:
self.assertEqual(m_result, fn_result)
if isinstance(b, float):
_helper(current_mobile_module_float, torch.Tensor.div_)
else:
_helper(current_mobile_module_int, torch.Tensor.div_)
def getScriptExportImportCopy(self, m, save_mobile_debug_info=True, also_test_file=False):
m_scripted = torch.jit.script(m)
if not also_test_file:
buffer = io.BytesIO(m_scripted._save_to_buffer_for_lite_interpreter(_save_mobile_debug_info=save_mobile_debug_info))
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
return mobile_module
with TemporaryFileName() as fname:
m_scripted._save_for_lite_interpreter(fname, _save_mobile_debug_info=save_mobile_debug_info)
mobile_module = _load_for_lite_interpreter(fname)
return mobile_module
def test_nest_typing_namedtuple_custom_classtype(self):
class Baz(NamedTuple):
di: torch.Tensor
class Foo(NamedTuple):
id: torch.Tensor
baz: Baz
class Bar(torch.nn.Module):
def __init__(self):
super(Bar, self).__init__()
self.foo = Foo(torch.tensor(1), Baz(torch.tensor(1)))
def forward(self, a: torch.Tensor):
self.foo = Foo(a, Baz(torch.tensor(1)))
return self.foo
sample_input = torch.tensor(5)
script_module = torch.jit.script(Bar())
script_module_result = script_module(sample_input)
buffer_mobile = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer_mobile.seek(0)
mobile_module = _load_for_lite_interpreter(buffer_mobile)
mobile_module_result = mobile_module(sample_input)
torch.testing.assert_allclose(script_module_result.baz.di,
mobile_module_result.baz.di)
def test_typing_dict_with_namedtuple(self):
class Foo(NamedTuple):
id: torch.Tensor
class Bar(torch.nn.Module):
def __init__(self):
super(Bar, self).__init__()
self.foo = Foo(torch.tensor(1))
def forward(self, a: torch.Tensor):
self.foo = Foo(a)
re: Dict[str, Foo] = dict()
re["test"] = Foo(a)
return self.foo, re["test"]
# The corresponding bytecode is
# (8,
# ('__torch__.___torch_mangle_2.Bar.forward',
# (('instructions',
# (('STOREN', 1, 2),
# ('DROPR', 1, 0),
# ('DICT_CONSTRUCT', 0, 0),
# ('STORE', 3, 0),
# ('LOAD', 3, 0),
# ('LOADC', 1, 0),
# ('MOVE', 2, 0),
# ('NAMED_TUPLE_CONSTRUCT', 1, 1),
# ('OP', 0, 0),
# ('MOVE', 3, 0),
# ('LOADC', 1, 0),
# ('DICT_INDEX', 0, 0),
# ('LOADC', 0, 0),
# ('TUPLE_INDEX', 0, 0),
# ('RET', 0, 0))),
# ('operators', (('aten::_set_item', 'str', 3),)),
# ('constants', (0, 'test')),
# ('types',
# ('Dict[str,__torch__.Foo[NamedTuple, [[id, Tensor]]]]',
# '__torch__.Foo[NamedTuple, [[id, Tensor]]]')),
# ('register_size', 3)),
# (('arguments',
# ((('name', 'self'),
# ('type', '__torch__.___torch_mangle_2.Bar'),
# ('default_value', None)),
# (('name', 'a'), ('type', 'Tensor'), ('default_value', None)))),
# ('returns',
# ((('name', ''), ('type', 'Tensor'), ('default_value', None)),)))))
sample_input = torch.tensor(5)
script_module = torch.jit.script(Bar())
script_module_result = script_module(sample_input)
buffer_mobile = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer_mobile.seek(0)
mobile_module = _load_for_lite_interpreter(buffer_mobile)
mobile_module_result = mobile_module(sample_input)
torch.testing.assert_allclose(script_module_result,
mobile_module_result)
def setUp(self):
super().setUp()
# Create Python and JIT versions of CompModule with lowered submodules.
compile_spec = {
"forward": {
"input_shapes": "((1, 1, 320, 240), (1, 3))",
"some_other_option": "True",
},
}
lowered_add = torch._C._jit_to_backend(
"backend_with_compiler_demo", torch.jit.script(BasicModuleAdd()),
compile_spec)
lowered_sub = torch._C._jit_to_backend(
"backend_with_compiler_demo",
torch.jit.script(CompModuleTestWithCompiler.BasicModuleSub()),
{"forward": {
"": ""
}})
self.module = CompModuleTestWithCompiler.CompModule(
lowered_add, lowered_sub)
self.scripted_module = torch.jit.script(
CompModuleTestWithCompiler.CompModule(lowered_add, lowered_sub))
# No backend version of CompModule currently, so this is filler.
self.lowered_module = self.scripted_module
# Create a mobile version of CompModule from JIT version
buffer = io.BytesIO(
self.scripted_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
self.mobile_module = _load_for_lite_interpreter(buffer)
def _compare_script_and_mobile(self, model: torch.nn.Module,
input: torch.Tensor):
qengine = "qnnpack"
with override_quantized_engine(qengine):
script_module = torch.jit.script(model)
script_module_result = script_module(input)
max_retry = 5
for retry in range(1, max_retry + 1):
# retires `max_retry` times; breaks iff succeeds else throws exception
try:
buffer = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
mobile_module_result = mobile_module(input)
torch.testing.assert_allclose(script_module_result,
mobile_module_result)
mobile_module_forward_result = mobile_module.forward(input)
torch.testing.assert_allclose(
script_module_result, mobile_module_forward_result)
mobile_module_run_method_result = mobile_module.run_method(
"forward", input)
torch.testing.assert_allclose(
script_module_result, mobile_module_run_method_result)
except AssertionError as e:
if retry == max_retry:
raise e
else:
continue
break
def test_backport_bytecode_from_file_to_buffer(self):
maximum_checked_in_model_version = max(
SCRIPT_MODULE_BYTECODE_PKL.keys())
script_module_v5_path = pytorch_test_dri / "cpp" / "jit" / SCRIPT_MODULE_BYTECODE_PKL[
maximum_checked_in_model_version]["model_name"]
if (maximum_checked_in_model_version > MINIMUM_TO_VERSION):
# Backport model to v4
script_module_v4_buffer = _backport_for_mobile_to_buffer(
script_module_v5_path, maximum_checked_in_model_version - 1)
buf = io.StringIO()
# Check version of the model v4 from backport
bytesio = io.BytesIO(script_module_v4_buffer)
backport_version = _get_model_bytecode_version(bytesio)
assert (backport_version == maximum_checked_in_model_version - 1)
# Load model v4 from backport and run forward method
bytesio = io.BytesIO(script_module_v4_buffer)
mobile_module = _load_for_lite_interpreter(bytesio)
module_input = 1
mobile_module_result = mobile_module(module_input)
expected_mobile_module_result = 3 * torch.ones([2, 4],
dtype=torch.float64)
torch.testing.assert_allclose(mobile_module_result,
expected_mobile_module_result)
def test_load_mobile_module(self):
class MyTestModule(torch.nn.Module):
def __init__(self):
super(MyTestModule, self).__init__()
def forward(self, x):
return x + 10
input = torch.tensor([1])
script_module = torch.jit.script(MyTestModule())
script_module_result = script_module(input)
buffer = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
mobile_module_result = mobile_module(input)
torch.testing.assert_close(script_module_result, mobile_module_result)
mobile_module_forward_result = mobile_module.forward(input)
torch.testing.assert_close(script_module_result,
mobile_module_forward_result)
mobile_module_run_method_result = mobile_module.run_method(
"forward", input)
torch.testing.assert_close(script_module_result,
mobile_module_run_method_result)
def _save_load_mobile_module(self, script_module: torch.jit.ScriptModule):
buffer = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter(
_save_mobile_debug_info=True))
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
return mobile_module
def test_versioned_logspace_out(self):
class Module(torch.nn.Module):
def __init__(self):
super(Module, self).__init__()
def forward(self, a: Union[int, float, complex],
b: Union[int, float, complex], out: torch.Tensor):
return torch.logspace(a, b, steps=100, out=out)
model_path = pytorch_test_dir + "/jit/fixtures/test_versioned_logspace_out_v8.ptl"
loaded_model = torch.jit.load(model_path)
buffer = io.BytesIO(
loaded_model._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
v8_mobile_module = _load_for_lite_interpreter(buffer)
current_mobile_module = self._save_load_mobile_module(Module)
sample_inputs = (
(3, 10, torch.empty((100, ), dtype=torch.int64),
torch.empty((100, ), dtype=torch.int64)),
(-10, 10, torch.empty((100, ), dtype=torch.int64),
torch.empty((100, ), dtype=torch.int64)),
(4.0, 6.0, torch.empty((100, ), dtype=torch.float64),
torch.empty((100, ), dtype=torch.float64)),
(3 + 4j, 4 + 5j, torch.empty((100, ), dtype=torch.complex64),
torch.empty((100, ), dtype=torch.complex64)),
)
for (start, end, out_for_old, out_for_new) in sample_inputs:
output = v8_mobile_module(start, end, out_for_old)
output_current = current_mobile_module(start, end, out_for_new)
# when no step is given, should have used 100
self.assertTrue(output.size(dim=0) == 100)
# "Upgraded" model should match the new version output
self.assertEqual(output, output_current)
def test_versioned_div_tensor(self):
def div_tensor_0_3(self, other):
if self.is_floating_point() or other.is_floating_point():
return self.true_divide(other)
return self.divide(other, rounding_mode='trunc')
model_path = pytorch_test_dir / "cpp" / "jit" / "upgrader_models" / "test_versioned_div_tensor_v2.ptl"
mobile_module_v2 = _load_for_lite_interpreter(str(model_path))
jit_module_v2 = torch.jit.load(str(model_path))
current_mobile_module = self._save_load_mobile_module(jit_module_v2)
vals = (2., 3., 2, 3)
for val_a, val_b in product(vals, vals):
a = torch.tensor((val_a, ))
b = torch.tensor((val_b, ))
def _helper(m, fn):
m_results = self._try_fn(m, a, b)
fn_result = self._try_fn(fn, a, b)
if isinstance(m_results, Exception):
self.assertTrue(isinstance(fn_result, Exception))
else:
for result in m_results:
print("result: ", result)
print("fn_result: ", fn_result)
print(result == fn_result)
self.assertTrue(result.eq(fn_result))
def test_backport_bytecode_from_file_to_file(self):
maximum_checked_in_model_version = max(SCRIPT_MODULE_BYTECODE_PKL.keys())
script_module_v5_path = pytorch_test_dir / "cpp" / "jit" / SCRIPT_MODULE_BYTECODE_PKL[
maximum_checked_in_model_version]["model_name"]
if (maximum_checked_in_model_version > MINIMUM_TO_VERSION):
with tempfile.TemporaryDirectory() as tmpdirname:
tmp_backport_model_path = Path(tmpdirname, "tmp_script_module_v5_backported_to_v4.ptl")
# backport from file
success = _backport_for_mobile(
script_module_v5_path,
tmp_backport_model_path,
maximum_checked_in_model_version - 1)
assert(success)
buf = io.StringIO()
torch.utils.show_pickle.main(
["", tmpdirname + "/" + tmp_backport_model_path.name + "@*/bytecode.pkl"],
output_stream=buf)
output = buf.getvalue()
expected_result = SCRIPT_MODULE_V4_BYTECODE_PKL
acutal_result_clean = "".join(output.split())
expect_result_clean = "".join(expected_result.split())
isMatch = fnmatch.fnmatch(acutal_result_clean, expect_result_clean)
assert(isMatch)
# Load model v4 and run forward method
mobile_module = _load_for_lite_interpreter(str(tmp_backport_model_path))
module_input = 1
mobile_module_result = mobile_module(module_input)
expected_mobile_module_result = 3 * torch.ones([2, 4], dtype=torch.float64)
torch.testing.assert_close(mobile_module_result, expected_mobile_module_result)
shutil.rmtree(tmpdirname)
def test_versioned_logspace(self):
class Module(torch.nn.Module):
def __init__(self):
super(Module, self).__init__()
def forward(self, a: Union[int, float, complex],
b: Union[int, float, complex]):
c = torch.logspace(a, b, steps=5)
d = torch.logspace(a, b, steps=100)
return c, d
scripted_module = torch.jit.load(
pytorch_test_dir + "/jit/fixtures/test_versioned_logspace_v8.ptl")
buffer = io.BytesIO(
scripted_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
v8_mobile_module = _load_for_lite_interpreter(buffer)
current_mobile_module = self._save_load_mobile_module(Module)
sample_inputs = ((3, 10), (-10, 10), (4.0, 6.0), (3 + 4j, 4 + 5j))
for (a, b) in sample_inputs:
(output_with_step, output_without_step) = v8_mobile_module(a, b)
(current_with_step,
current_without_step) = current_mobile_module(a, b)
# when no step is given, should have used 100
self.assertTrue(output_without_step.size(dim=0) == 100)
self.assertTrue(output_with_step.size(dim=0) == 5)
# outputs should be equal to the newest version
self.assertEqual(output_with_step, current_with_step)
self.assertEqual(output_without_step, current_without_step)
def test_versioned_div_scalar_scalar(self):
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, a: float, b: int, c: float, d: int):
result_0 = a / b
result_1 = a / c
result_2 = b / c
result_3 = b / d
return (result_0, result_1, result_2, result_3)
try:
v3_mobile_module = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_scalar_scalar_v2.ptl"
)
except Exception as e:
self.skipTest("Failed to load fixture!")
current_mobile_module = self._save_load_mobile_module(MyModule)
def _helper(m, fn):
vals = (5., 3, 2., 7)
m_result = m(*vals)
fn_result = fn(*vals)
for mr, hr in zip(m_result, fn_result):
self.assertEqual(mr, hr)
_helper(v3_mobile_module, current_mobile_module)
def test_versioned_div_tensor_out(self):
def historic_div_out(self, other, out):
if self.is_floating_point() or other.is_floating_point(
) or out.is_floating_point():
return torch.true_divide(self, other, out=out)
return torch.divide(self, other, out=out, rounding_mode='trunc')
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, a, b, out):
return a.div(b, out=out)
try:
v3_module = torch.jit.load(
pytorch_test_dir +
"/jit/fixtures/test_versioned_div_tensor_out_v3.pt")
v3_mobile_module = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_tensor_out_v2.ptl"
)
except Exception as e:
self.skipTest("Failed to load fixture!")
self._verify_count("aten::div", v3_module,
2) # true_divide and divide alias to div
self._verify_count('prim::Constant[value="trunc"]', v3_module,
1) # rounding_mode argument
current_module = self._save_load_module(MyModule)
current_mobile_module = self._save_load_mobile_module(MyModule)
self._verify_count("aten::div", current_module, 1)
vals = (2., 3., 2, 3)
for val_a, val_b in product(vals, vals):
a = torch.tensor((val_a, ))
b = torch.tensor((val_b, ))
for out in (torch.empty((1, )), torch.empty((1, ),
dtype=torch.long)):
def _helper(m, fn):
fn_result = None
if fn is torch.div:
fn_result = self._try_fn(fn, a, b, out=out.clone())
else:
fn_result = self._try_fn(fn, a, b, out.clone())
m_result = self._try_fn(m, a, b, out)
if isinstance(m_result, Exception):
self.assertTrue(fn_result, Exception)
else:
self.assertEqual(m_result, fn_result)
self.assertEqual(m_result, out)
_helper(v3_module, historic_div_out)
_helper(current_module, torch.div)
_helper(v3_mobile_module, historic_div_out)
_helper(current_mobile_module, torch.div)
def setUp(self):
super().setUp()
self.module = CompModule()
self.scripted_module = torch.jit.script(self.module)
buffer = io.BytesIO(self.scripted_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
self.mobile_module = _load_for_lite_interpreter(buffer)
def test_versioned_div_tensor(self):
def historic_div(self, other):
if self.is_floating_point() or other.is_floating_point():
return self.true_divide(other)
return self.divide(other, rounding_mode='trunc')
# Tensor x Tensor
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, a, b):
result_0 = a / b
result_1 = torch.div(a, b)
result_2 = a.div(b)
return result_0, result_1, result_2
# Loads historic module
try:
v3_module = torch.jit.load(
pytorch_test_dir +
"/jit/fixtures/test_versioned_div_tensor_v3.pt")
v3_mobile_module = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_tensor_v2.ptl")
except Exception as e:
self.skipTest("Failed to load fixture!")
self._verify_count("aten::div", v3_module,
6) # true_divide and divide alias to div
self._verify_count('prim::Constant[value="trunc"]', v3_module,
1) # rounding_mode argument
current_module = self._save_load_module(MyModule)
current_mobile_module = self._save_load_mobile_module(MyModule)
self._verify_count("aten::div", current_module, 3)
vals = (2., 3., 2, 3)
for val_a, val_b in product(vals, vals):
a = torch.tensor((val_a, ))
b = torch.tensor((val_b, ))
def _helper(m, fn):
m_results = self._try_fn(m, a, b)
fn_result = self._try_fn(fn, a, b)
if isinstance(m_results, Exception):
self.assertTrue(isinstance(fn_result, Exception))
else:
for result in m_results:
self.assertEqual(result, fn_result)
_helper(v3_module, historic_div)
_helper(v3_mobile_module, historic_div)
_helper(current_module, torch.div)
_helper(current_mobile_module, torch.div)
def test_all_backport_functions(self):
# Backport from the latest bytecode version to the minimum support version
# Load, run the backport model, and check version
class TestModule(torch.nn.Module):
def __init__(self, v):
super().__init__()
self.x = v
def forward(self, y: int):
increment = torch.ones([2, 4], dtype=torch.float64)
return self.x + y + increment
module_input = 1
expected_mobile_module_result = 3 * torch.ones([2, 4],
dtype=torch.float64)
# temporary input model file and output model file will be exported in the temporary folder
with tempfile.TemporaryDirectory() as tmpdirname:
tmp_input_model_path = Path(tmpdirname, "tmp_script_module.ptl")
script_module = torch.jit.script(TestModule(1))
optimized_scripted_module = optimize_for_mobile(script_module)
exported_optimized_scripted_module = optimized_scripted_module._save_for_lite_interpreter(
str(tmp_input_model_path))
current_from_version = _get_model_bytecode_version(
tmp_input_model_path)
current_to_version = current_from_version - 1
tmp_output_model_path = Path(tmpdirname,
"tmp_script_module_backport.ptl")
while current_to_version >= MINIMUM_TO_VERSION:
# Backport the latest model to `to_version` to a tmp file "tmp_script_module_backport"
backport_success = _backport_for_mobile(
tmp_input_model_path, tmp_output_model_path,
current_to_version)
assert (backport_success)
backport_version = _get_model_bytecode_version(
tmp_output_model_path)
assert (backport_version == current_to_version)
# Load model and run forward method
mobile_module = _load_for_lite_interpreter(
str(tmp_input_model_path))
mobile_module_result = mobile_module(module_input)
torch.testing.assert_allclose(mobile_module_result,
expected_mobile_module_result)
current_to_version -= 1
# Check backport failure case
backport_success = _backport_for_mobile(tmp_input_model_path,
tmp_output_model_path,
MINIMUM_TO_VERSION - 1)
assert (not backport_success)
# need to clean the folder before it closes, otherwise will run into git not clean error
shutil.rmtree(tmpdirname)
def test_versioned_div_tensor_inplace(self):
def historic_div_(self, other):
if self.is_floating_point() or other.is_floating_point():
return self.true_divide_(other)
return self.divide_(other, rounding_mode='trunc')
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, a, b):
a /= b
return a
try:
v3_module = torch.jit.load(
pytorch_test_dir +
"/jit/fixtures/test_versioned_div_tensor_inplace_v3.pt")
v3_mobile_module = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_tensor_inplace_v2.ptl"
)
except Exception as e:
self.skipTest("Failed to load fixture!")
self._verify_count("aten::div", v3_module,
2) # true_divide and divide both alias to div
self._verify_count('prim::Constant[value="trunc"]', v3_module,
1) # rounding_mode argument
current_module = self._save_load_module(MyModule)
current_mobile_module = self._save_load_mobile_module(MyModule)
self._verify_count("aten::div", current_module, 1)
vals = (2., 3., 2, 3)
for val_a, val_b in product(vals, vals):
a = torch.tensor((val_a, ))
b = torch.tensor((val_b, ))
def _helper(m, fn):
fn_result = self._try_fn(fn, a.clone(), b)
m_result = self._try_fn(m, a, b)
if isinstance(m_result, Exception):
self.assertTrue(fn_result, Exception)
else:
self.assertEqual(m_result, fn_result)
self.assertEqual(m_result, a)
_helper(v3_module, historic_div_)
_helper(v3_mobile_module, historic_div_)
# Recreates a since it was modified in place
a = torch.tensor((val_a, ))
_helper(current_module, torch.Tensor.div_)
_helper(current_mobile_module, torch.Tensor.div_)
def runModule(module):
buffer = io.BytesIO(module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
lite_module = _load_for_lite_interpreter(buffer)
if lite_module.find_method("get_all_bundled_inputs"):
# run with the first bundled input
input = lite_module.run_method("get_all_bundled_inputs")[0]
lite_module.forward(*input)
else:
# assuming model has no input
lite_module()
def wrapper_save():
_MODELS.append(cls)
model = cls()
scripted = torch.jit.script(model)
buffer = BytesIO(scripted._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
ops = _export_operator_list(mobile_module)
_OPERATORS.update(ops)
path = f"./{cls.__name__}.ptl"
_FILENAMES.append(path)
scripted._save_for_lite_interpreter(path)
def test_stacktrace_interface_call(self):
@torch.jit.interface
class Forward(torch.nn.Module):
def forward(self, x) -> torch.Tensor:
pass
def forwardError(self, x) -> torch.Tensor:
pass
class B(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return x
def forwardError(self, x):
return self.call() + x
def call(self):
return torch.ones(-1)
class A(torch.nn.Module):
b: Forward
def __init__(self):
super().__init__()
self.b = B()
def forward(self):
self.b.forward(torch.ones(1))
self.b.forwardError(torch.ones(1))
a = torch.jit.script(A())
torch._C._enable_mobile_interface_call_export()
buffer = io.BytesIO(
a._save_to_buffer_for_lite_interpreter(
_save_mobile_debug_info=True))
buffer.seek(0)
mobile_module = _load_for_lite_interpreter(buffer)
try:
mobile_module()
self.assertTrue(False)
except RuntimeError as exp:
FileCheck().check("Trying to create tensor with negative dimension") \
.check("Traceback of TorchScript") \
.check("self.b.forwardError").check_next("~~~~~~~~~~~~~~~~~~~ <--- HERE") \
.check("return self.call").check_next("~~~~~~~~~ <--- HERE") \
.check("return torch.ones").check_next("~~~~~~~~~~ <--- HERE").run(str(exp))
def test_save_mobile_module_with_debug_info_with_trace(self):
class A(torch.nn.Module):
def __init__(self):
super(A, self).__init__()
def forward(self, x, y):
return x * y
class B(torch.nn.Module):
def __init__(self):
super(B, self).__init__()
self.A0 = A()
self.A1 = A()
def forward(self, x, y, z):
return self.A0(x, y) + self.A1(y, z)
for export_method in ['trace', 'script']:
x = torch.rand((2, 3))
y = torch.rand((2, 3))
z = torch.rand((2, 3))
if export_method == 'trace':
trace_module = torch.jit.trace(B(), [x, y, z])
else:
trace_module = torch.jit.script(B())
exported_module = trace_module._save_to_buffer_for_lite_interpreter(
_save_mobile_debug_info=True)
buffer = io.BytesIO(exported_module)
buffer.seek(0)
assert (b"callstack_debug_map.pkl" in exported_module)
mobile_module = _load_for_lite_interpreter(buffer)
with self.assertRaisesRegex(
RuntimeError,
r"Module hierarchy:top\(B\)::<unknown>.A0\(A\)::forward.aten::mul"
):
x = torch.rand((2, 3))
y = torch.rand((8, 10))
z = torch.rand((8, 10))
mobile_module(x, y, z)
with self.assertRaisesRegex(
RuntimeError,
r"Module hierarchy:top\(B\)::<unknown>.A1\(A\)::forward.aten::mul"
):
x = torch.rand((2, 3))
y = torch.rand((2, 3))
z = torch.rand((8, 10))
mobile_module(x, y, z)
def test_versioned_div_tensor_out(self, sample_input):
def historic_div_out(self, other, out):
if self.is_floating_point() or other.is_floating_point(
) or out.is_floating_point():
return torch.true_divide(self, other, out=out)
return torch.divide(self, other, out=out, rounding_mode='trunc')
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, a, b, out):
return a.div(b, out=out)
try:
v3_mobile_module = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_tensor_out_v2.ptl"
)
except Exception as e:
self.skipTest("Failed to load fixture!")
current_mobile_module = self._save_load_mobile_module(MyModule)
for val_a, val_b in product(sample_input, sample_input):
a = torch.tensor((val_a, ))
b = torch.tensor((val_b, ))
for out in (torch.empty((1, )), torch.empty((1, ),
dtype=torch.long)):
def _helper(m, fn):
fn_result = None
if fn is torch.div:
fn_result = self._try_fn(fn, a, b, out=out.clone())
else:
fn_result = self._try_fn(fn, a, b, out.clone())
m_result = self._try_fn(m, a, b, out)
if isinstance(m_result, Exception):
self.assertTrue(fn_result, Exception)
else:
self.assertEqual(m_result, fn_result)
self.assertEqual(m_result, out)
_helper(v3_mobile_module, historic_div_out)
_helper(current_mobile_module, torch.div)
def test_return_collections_namedtuple(self):
myNamedTuple = namedtuple('myNamedTuple', [('a')])
class MyTestModule(torch.nn.Module):
def forward(self, a: torch.Tensor):
return myNamedTuple(a)
sample_input = torch.Tensor(1)
script_module = torch.jit.script(MyTestModule())
script_module_result = script_module(sample_input)
buffer_mobile = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter())
buffer_mobile.seek(0)
mobile_module = _load_for_lite_interpreter(buffer_mobile)
mobile_module_result = mobile_module(sample_input)
torch.testing.assert_allclose(script_module_result,
mobile_module_result)
def test_method_calls_with_optional_arg(self):
class A(torch.nn.Module):
def __init__(self):
super(A, self).__init__()
# opt arg in script-to-script invocation
def forward(self, x, two: int = 2):
return x + two
class B(torch.nn.Module):
def __init__(self):
super(B, self).__init__()
self.A0 = A()
# opt arg in Python-to-script invocation
def forward(self, x, one: int = 1):
return self.A0(x) + one
script_module = torch.jit.script(B())
buffer = io.BytesIO(
script_module._save_to_buffer_for_lite_interpreter()
)
mobile_module = _load_for_lite_interpreter(buffer)
input = torch.tensor([5])
script_module_forward_result = script_module.forward(input)
mobile_module_forward_result = mobile_module.forward(input)
torch.testing.assert_close(
script_module_forward_result,
mobile_module_forward_result
)
# change ref only
script_module_forward_result = script_module.forward(input, 2)
self.assertFalse(
(script_module_forward_result == mobile_module_forward_result)
.all()
.item()
)
# now both match again
mobile_module_forward_result = mobile_module.forward(input, 2)
torch.testing.assert_close(
script_module_forward_result,
mobile_module_forward_result
)
def setUp(self):
super().setUp()
# Create Python, JIT and backend versions of BasicModuleAdd.
self.module = BasicModuleAdd()
self.scripted_module = torch.jit.script(BasicModuleAdd())
compile_spec = {
"forward": {
"input_shapes": "((1, 1, 320, 240), (1, 3))",
"some_other_option": "True",
},
}
self.lowered_module = torch._C._jit_to_backend(
"backend_with_compiler_demo", self.scripted_module, compile_spec)
# Create mobile version of BasicModuleAdd
buffer = io.BytesIO(self.lowered_module._save_to_buffer_for_lite_interpreter())
buffer.seek(0)
self.mobile_module = _load_for_lite_interpreter(buffer)
def test_versioned_div_tensor(self, sample_input):
def historic_div(self, other):
if self.is_floating_point() or other.is_floating_point():
return self.true_divide(other)
return self.divide(other, rounding_mode='trunc')
# Tensor x Tensor
class MyModule(torch.nn.Module):
def __init__(self):
super(MyModule, self).__init__()
def forward(self, a, b):
result_0 = a / b
result_1 = torch.div(a, b)
result_2 = a.div(b)
return result_0, result_1, result_2
# Loads historic module
try:
v3_mobile_module = _load_for_lite_interpreter(
pytorch_test_dir +
"/cpp/jit/upgrader_models/test_versioned_div_tensor_v2.ptl")
except Exception as e:
self.skipTest("Failed to load fixture!")
current_mobile_module = self._save_load_mobile_module(MyModule)
for val_a, val_b in product(sample_input, sample_input):
a = torch.tensor((val_a, ))
b = torch.tensor((val_b, ))
def _helper(m, fn):
m_results = self._try_fn(m, a, b)
fn_result = self._try_fn(fn, a, b)
if isinstance(m_results, Exception):
self.assertTrue(isinstance(fn_result, Exception))
else:
for result in m_results:
self.assertEqual(result, fn_result)
_helper(v3_mobile_module, historic_div)
_helper(current_mobile_module, torch.div)
def test_bundled_input_with_dynamic_type(self):
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(
self,
x: Dict[int, torch.Tensor],
y: Dict[int, torch.Tensor],
z: Dict[int, torch.Tensor],
):
return x
model = Model()
script_module = torch.jit.script(model)
sample_input = {
script_module.forward: [
(
{0: torch.ones(1)},
{1: torch.ones(1)},
{2: torch.ones(1)},
)
]
}
bundled_model = torch.utils.bundled_inputs.bundle_inputs(
script_module, sample_input
)
buf = bundled_model._save_to_buffer_for_lite_interpreter()
mobile_module = _load_for_lite_interpreter(io.BytesIO(buf))
i = mobile_module.run_method("get_all_bundled_inputs")
self.assertEqual(
i[0],
(
{0: torch.ones(1)},
{1: torch.ones(1)},
{2: torch.ones(1)},
),
)
