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_package_fx_custom_tracer(self):
from package_a.test_all_leaf_modules_tracer import TestAllLeafModulesTracer
from package_a.test_module import SimpleTest, ModWithTwoSubmodsAndTensor
class SpecialGraphModule(torch.fx.GraphModule):
def __init__(self, root, graph, info):
super().__init__(root, graph)
self.info = info
sub_module = SimpleTest()
module = ModWithTwoSubmodsAndTensor(
torch.ones(3),
sub_module,
sub_module,
)
tracer = TestAllLeafModulesTracer()
graph = tracer.trace(module)
self.assertEqual(graph._tracer_cls, TestAllLeafModulesTracer)
gm = SpecialGraphModule(module, graph, "secret")
self.assertEqual(gm._tracer_cls, TestAllLeafModulesTracer)
f = BytesIO()
with PackageExporter(f) as pe:
pe.intern("**")
pe.save_pickle("model", "model.pkl", gm)
f.seek(0)
pi = PackageImporter(f)
loaded_gm = pi.load_pickle("model", "model.pkl")
self.assertEqual(
type(loaded_gm).__class__.__name__, SpecialGraphModule.__class__.__name__
)
self.assertEqual(loaded_gm.info, "secret")
input_x = torch.randn(3)
self.assertEqual(loaded_gm(input_x), gm(input_x))