def test_scriptmodules_repeat_save(self):
"""
Test to verify saving and loading same ScriptModule object works
across multiple packages.
"""
from package_a.test_module import ModWithSubmodAndTensor, ModWithTensor
scripted_mod_0 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
scripted_mod_1 = torch.jit.script(
ModWithSubmodAndTensor(torch.rand(1, 2, 3),
ModWithTensor(torch.rand(1, 2, 3))))
buffer_0 = BytesIO()
with PackageExporter(buffer_0) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_0)
buffer_0.seek(0)
importer_0 = PackageImporter(buffer_0)
loaded_module_0 = importer_0.load_pickle("res", "mod1.pkl")
buffer_1 = BytesIO()
with PackageExporter(buffer_1) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_1)
e.save_pickle("res", "mod2.pkl", loaded_module_0)
buffer_1.seek(0)
importer_1 = PackageImporter(buffer_1)
loaded_module_1 = importer_1.load_pickle("res", "mod1.pkl")
reloaded_module_0 = importer_1.load_pickle("res", "mod2.pkl")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_module_0(input), scripted_mod_0(input))
self.assertEqual(loaded_module_0(input), reloaded_module_0(input))
self.assertEqual(loaded_module_1(input), scripted_mod_1(input))
def test_load_shared_tensors_repackaged(self):
"""
Test tensors shared across eager and ScriptModules on load
are the same across multiple package saves and loads. This is
an important test because not all of the tensor information is restored
in python between packages. The python identity is not maintained, but
the backing cpp TensorImpl is. We load/save storages based off of this
cpp TensorImpl and not the python identity.
"""
from package_a.test_module import (
ModWithTensor,
ModWithTwoSubmodsAndTensor,
)
shared_tensor = torch.ones(3, 3)
scripted_mod_0 = torch.jit.script(ModWithTensor(shared_tensor))
scripted_mod_1 = torch.jit.script(ModWithTensor(shared_tensor))
mod1 = ModWithTwoSubmodsAndTensor(shared_tensor, scripted_mod_0, scripted_mod_1)
buffer_0 = BytesIO()
with PackageExporter(buffer_0) as e:
e.intern("**")
e.save_pickle("res", "mod1.pkl", mod1)
buffer_0.seek(0)
importer_0 = PackageImporter(buffer_0)
loaded_mod_0 = importer_0.load_pickle("res", "mod1.pkl")
buffer_1 = BytesIO()
with PackageExporter(buffer_1, importer=importer_0) as e:
e.intern("**")
e.save_pickle("res", "mod1.pkl", loaded_mod_0)
buffer_1.seek(0)
importer = PackageImporter(buffer_1)
loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
self.assertEqual(
loaded_mod_1.tensor.storage()._cdata,
loaded_mod_1.sub_mod_0.tensor.storage()._cdata,
)
self.assertEqual(
loaded_mod_1.tensor.storage()._cdata,
loaded_mod_1.sub_mod_1.tensor.storage()._cdata,
)
loaded_mod_1.tensor.add_(
torch.ones(3, 3)
) # all tensors should reflect this change
self.assertTrue(
torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_0.tensor)
)
self.assertTrue(
torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_1.tensor)
)
def test_load_shared_tensors(self):
"""
Test tensors shared across eager and ScriptModules on load
are the same.
"""
from package_a.test_module import (
ModWithTensor,
ModWithTwoSubmodsAndTensor,
)
shared_tensor = torch.ones(3, 3)
scripted_mod_0 = torch.jit.script(ModWithTensor(shared_tensor))
scripted_mod_1 = torch.jit.script(ModWithTensor(shared_tensor))
mod1 = ModWithTwoSubmodsAndTensor(shared_tensor, scripted_mod_0, scripted_mod_1)
self.assertEqual(
shared_tensor.storage()._cdata,
scripted_mod_0.tensor.storage()._cdata,
)
self.assertEqual(
shared_tensor.storage()._cdata,
scripted_mod_1.tensor.storage()._cdata,
)
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.intern("**")
e.save_pickle("res", "mod1.pkl", mod1)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
self.assertEqual(
loaded_mod_1.tensor.storage()._cdata,
loaded_mod_1.sub_mod_0.tensor.storage()._cdata,
)
self.assertEqual(
loaded_mod_1.tensor.storage()._cdata,
loaded_mod_1.sub_mod_1.tensor.storage()._cdata,
)
loaded_mod_1.tensor.add_(torch.ones(3, 3))
self.assertTrue(
torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_0.tensor)
)
self.assertTrue(
torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_1.tensor)
)
def test_save_shared_tensors(self):
"""
Test tensors shared across eager and ScriptModules are serialized once.
"""
from package_a.test_module import ModWithSubmodAndTensor, ModWithTensor
shared_tensor = torch.rand(2, 3, 4)
scripted_mod = torch.jit.script(ModWithTensor(shared_tensor))
mod1 = ModWithSubmodAndTensor(shared_tensor, scripted_mod)
mod2 = ModWithSubmodAndTensor(shared_tensor, scripted_mod)
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.intern("**")
e.save_pickle("res", "tensor", shared_tensor)
e.save_pickle("res", "mod1.pkl", mod1)
e.save_pickle("res", "mod2.pkl", mod2)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
# assert that there is only one storage stored in package
file_structure = importer.file_structure(include=".data/*.storage")
self.assertTrue(len(file_structure.children[".data"].children) == 1)
input = torch.rand(2, 3, 4)
self.assertEqual(loaded_mod_1(input), mod1(input))
def test_scriptobject_failure_message(self):
"""
Test basic saving and loading of a ScriptModule in a directory.
Currently not supported.
"""
from package_a.test_module import ModWithTensor
scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
filename = self.temp()
with PackageExporter(filename, verbose=False) as e:
e.save_pickle("res", "mod.pkl", scripted_mod)
zip_file = zipfile.ZipFile(filename, "r")
with self.assertRaisesRegex(
RuntimeError,
"Loading ScriptObjects from a PackageImporter created from a "
"directory is not supported. Use a package archive file instead.",
):
with TemporaryDirectory() as temp_dir:
zip_file.extractall(path=temp_dir)
dir_importer = PackageImporter(
Path(temp_dir) / Path(filename).name)
dir_mod = dir_importer.load_pickle("res", "mod.pkl")
def test_save_repeat_scriptmodules(self):
"""
Test to verify saving multiple different modules and
repeats of same scriptmodule in package works. Also tests that
PyTorchStreamReader isn't having code hidden from
PyTorchStreamWriter writing ScriptModule code files multiple times.
"""
from package_a.test_module import (
ModWithSubmodAndTensor,
ModWithTensor,
SimpleTest,
)
scripted_mod_0 = torch.jit.script(SimpleTest())
scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
scripted_mod_2 = torch.jit.script(
ModWithSubmodAndTensor(
torch.rand(1, 2, 3), ModWithTensor(torch.rand(1, 2, 3))
)
)
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.save_pickle("res", "mod0.pkl", scripted_mod_0)
e.save_pickle("res", "mod1.pkl", scripted_mod_1)
e.save_pickle("res", "mod2.pkl", scripted_mod_0)
e.save_pickle("res", "mod3.pkl", scripted_mod_1)
e.save_pickle("res", "mod4.pkl", scripted_mod_2)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod_0 = importer.load_pickle("res", "mod0.pkl")
loaded_mod_1 = importer.load_pickle("res", "mod3.pkl")
loaded_mod_2 = importer.load_pickle("res", "mod4.pkl")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod_0(input), scripted_mod_0(input))
self.assertEqual(loaded_mod_1(input), scripted_mod_1(input))
self.assertEqual(loaded_mod_2(input), scripted_mod_2(input))
def test_save_scriptmodule_file(self):
"""
Test basic saving of ScriptModule in file.
"""
from package_a.test_module import ModWithTensor
scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
filename = self.temp()
with PackageExporter(filename) as e:
e.save_pickle("res", "mod.pkl", scripted_mod)
importer = PackageImporter(filename)
loaded_mod = importer.load_pickle("res", "mod.pkl")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod(input), scripted_mod(input))
def test_save_scriptmodule(self):
"""
Test basic saving of ScriptModule.
"""
from package_a.test_module import ModWithTensor
scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.save_pickle("res", "mod.pkl", scripted_mod)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod = importer.load_pickle("res", "mod.pkl", map_location="cpu")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod(input), scripted_mod(input))
def test_save_independent_scriptmodules(self):
"""
Test to verify saving multiple ScriptModules with completely
separate code works.
"""
from package_a.test_module import ModWithTensor, SimpleTest
scripted_mod_0 = torch.jit.script(SimpleTest())
scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_0)
e.save_pickle("res", "mod2.pkl", scripted_mod_1)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod_0 = importer.load_pickle("res", "mod1.pkl")
loaded_mod_1 = importer.load_pickle("res", "mod2.pkl")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod_0(input), scripted_mod_0(input))
self.assertEqual(loaded_mod_1(input), scripted_mod_1(input))
def test_save_scriptmodules_in_container(self):
"""
Test saving of ScriptModules inside of container. Checks that relations
between shared modules are upheld.
"""
from package_a.test_module import ModWithSubmodAndTensor, ModWithTensor
scripted_mod_a = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
scripted_mod_b = torch.jit.script(
ModWithSubmodAndTensor(torch.rand(1, 2, 3), scripted_mod_a))
script_mods_list = [scripted_mod_a, scripted_mod_b]
buffer = BytesIO()
with PackageExporter(buffer) as e:
e.save_pickle("res", "list.pkl", script_mods_list)
buffer.seek(0)
importer = PackageImporter(buffer)
loaded_mod_list = importer.load_pickle("res", "list.pkl")
input = torch.rand(1, 2, 3)
self.assertEqual(loaded_mod_list[0](input), scripted_mod_a(input))
self.assertEqual(loaded_mod_list[1](input), scripted_mod_b(input))
def test_save_scriptmodule_only_necessary_code(self):
"""
Test to verify when saving multiple packages with same CU
that packages don't include unnecessary torchscript code files.
The TorchVision code should only be saved in the package that
relies on it.
"""
from package_a.test_module import ModWithTensor
class ModWithTorchVision(torch.nn.Module):
def __init__(self, name: str):
super().__init__()
self.tvmod = resnet18()
def forward(self, input):
return input * 4
scripted_mod_0 = torch.jit.script(ModWithTorchVision("foo"))
scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
buffer_0 = BytesIO()
with PackageExporter(buffer_0) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_0)
buffer_0.seek(0)
importer_0 = importer = PackageImporter(buffer_0)
buffer_1 = BytesIO()
with PackageExporter(buffer_1) as e:
e.save_pickle("res", "mod1.pkl", scripted_mod_1)
buffer_1.seek(0)
importer_1 = PackageImporter(buffer_1)
self.assertTrue("torchvision" in str(importer_0.file_structure()))
self.assertFalse("torchvision" in str(importer_1.file_structure()))
