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()))