mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
2a7a39c1af
Summary: This PR updates `torch::pickle_save` to use the new zipfile format introduced in #29232 and adds `torch::pickle_load` which can decode the zipfile format. Now that `torch.save/load` use this format as well (if the `_use_new_zipfile_serialization` flag is `True`), raw values saved in Python can be loaded in C++ and vice versa. Fixes #20356 ](https://our.intern.facebook.com/intern/diff/18607087/) Pull Request resolved: https://github.com/pytorch/pytorch/pull/30108 Pulled By: driazati Differential Revision: D18607087 fbshipit-source-id: 067cdd5b1cf9c30ddc7e2e5021a8cceee62d8a14
75 lines
1.4 KiB
Python
75 lines
1.4 KiB
Python
import sys
|
|
import os
|
|
import torch
|
|
|
|
|
|
class Setup(object):
|
|
def setup(self):
|
|
raise NotImplementedError()
|
|
|
|
def shutdown(self):
|
|
raise NotImplementedError()
|
|
|
|
|
|
class FileSetup(object):
|
|
path = None
|
|
|
|
def shutdown(self):
|
|
if os.path.exists(self.path):
|
|
os.remove(self.path)
|
|
pass
|
|
|
|
|
|
class EvalModeForLoadedModule(FileSetup):
|
|
path = 'dropout_model.pt'
|
|
|
|
def setup(self):
|
|
class Model(torch.jit.ScriptModule):
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
self.dropout = torch.nn.Dropout(0.1)
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, x):
|
|
x = self.dropout(x)
|
|
return x
|
|
|
|
model = Model()
|
|
model = model.train()
|
|
model.save(self.path)
|
|
|
|
|
|
class SerializationInterop(FileSetup):
|
|
path = 'ivalue.pt'
|
|
|
|
def setup(self):
|
|
ones = torch.ones(2, 2)
|
|
twos = torch.ones(3, 5) * 2
|
|
|
|
value = (ones, twos)
|
|
|
|
torch.save(value, self.path, _use_new_zipfile_serialization=True)
|
|
|
|
|
|
tests = [
|
|
EvalModeForLoadedModule(),
|
|
SerializationInterop(),
|
|
]
|
|
|
|
def setup():
|
|
for test in tests:
|
|
test.setup()
|
|
|
|
|
|
def shutdown():
|
|
for test in tests:
|
|
test.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
command = sys.argv[1]
|
|
if command == "setup":
|
|
setup()
|
|
elif command == "shutdown":
|
|
shutdown()
|