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_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 check_model_version(model_path, expect_version):
actual_version = _get_model_bytecode_version(model_path)
assert(actual_version == expect_version)
def test_get_model_bytecode_version(self):
script_module_v4 = pytorch_test_dri / "cpp" / "jit" / "script_module_v4.ptl"
version_v4 = _get_model_bytecode_version(script_module_v4)
assert (version_v4 == 4)