def load(cls, path):
with tarfile.open(utils.ensure_ext(path, 'tar'), 'r') as tar:
commit = utils.get_gzip_from_tar(tar, 'pie-commit.zip')
if pie.__commit__ != commit:
logging.warn(
("Model {} was serialized with a previous "
"version of `pie`. This might result in issues. "
"Model commit is {}, whereas current `pie` commit is {}."
).format(path, commit, pie.__commit__))
# load label encoder
le = pie.dataset.MultiLabelEncoder.load_from_string(
utils.get_gzip_from_tar(tar, 'label_encoder.zip'))
# load model parameters
params = json.loads(utils.get_gzip_from_tar(tar, 'parameters.zip'))
# instantiate model
model = Encoder(le, *params['args'], **params['kwargs'])
# load state_dict
with utils.tmpfile() as tmppath:
tar.extract('state_dict.pt', path=tmppath)
dictpath = os.path.join(tmppath, 'state_dict.pt')
model.load_state_dict(torch.load(dictpath, map_location='cpu'))
model.eval()
return model
def save(self, path):
path = utils.ensure_ext(path, 'tar')
# create dir if needed
dirname = os.path.dirname(path)
if dirname and not os.path.isdir(dirname):
os.makedirs(dirname)
with tarfile.open(path, 'w') as tar:
# serialize label_encoder
string, path = json.dumps(
self.label_encoder.jsonify()), 'label_encoder.zip'
utils.add_gzip_to_tar(string, path, tar)
# serialize parameters
string, path = json.dumps(
self.get_args_and_kwargs()), 'parameters.zip'
utils.add_gzip_to_tar(string, path, tar)
# serialize weights
with utils.tmpfile() as tmppath:
torch.save(self.state_dict(), tmppath)
tar.add(tmppath, arcname='state_dict.pt')
# serialize current pie commit
string, path = pie.__commit__, 'pie-commit.zip'
utils.add_gzip_to_tar(string, path, tar)
def load(fpath):
"""
Load model from path
"""
import pie
with tarfile.open(utils.ensure_ext(fpath, 'tar'), 'r') as tar:
# check commit
try:
commit = utils.get_gzip_from_tar(tar, 'pie-commit.zip')
except Exception:
commit = None
if (pie.__commit__ and commit) and pie.__commit__ != commit:
logging.warn(
("Model {} was serialized with a previous "
"version of `pie`. This might result in issues. "
"Model commit is {}, whereas current `pie` commit is {}."
).format(fpath, commit, pie.__commit__))
# load label encoder
le = MultiLabelEncoder.load_from_string(
utils.get_gzip_from_tar(tar, 'label_encoder.zip'))
# load tasks
tasks = json.loads(utils.get_gzip_from_tar(tar, 'tasks.zip'))
# load model parameters
params = json.loads(utils.get_gzip_from_tar(tar, 'parameters.zip'))
# instantiate model
model_type = getattr(pie.models,
utils.get_gzip_from_tar(tar, 'class.zip'))
with utils.shutup():
model = model_type(le, tasks, *params['args'],
**params['kwargs'])
# load settings
try:
settings = Settings(
json.loads(utils.get_gzip_from_tar(tar, 'settings.zip')))
model._settings = settings
except Exception:
logging.warn(
"Couldn't load settings for model {}!".format(fpath))
# load state_dict
with utils.tmpfile() as tmppath:
tar.extract('state_dict.pt', path=tmppath)
dictpath = os.path.join(tmppath, 'state_dict.pt')
model.load_state_dict(torch.load(dictpath, map_location='cpu'))
model.eval()
return model
def save(self, fpath, infix=None, settings=None):
"""
Serialize model to path
"""
import pie
fpath = utils.ensure_ext(fpath, 'tar', infix)
# create dir if necessary
dirname = os.path.dirname(fpath)
if not os.path.isdir(dirname):
os.makedirs(dirname)
with tarfile.open(fpath, 'w') as tar:
# serialize label_encoder
string = json.dumps(self.label_encoder.jsonify())
path = 'label_encoder.zip'
utils.add_gzip_to_tar(string, path, tar)
# serialize tasks
string, path = json.dumps(self.tasks), 'tasks.zip'
utils.add_gzip_to_tar(string, path, tar)
# serialize model class
string, path = str(type(self).__name__), 'class.zip'
utils.add_gzip_to_tar(string, path, tar)
# serialize parameters
string, path = json.dumps(
self.get_args_and_kwargs()), 'parameters.zip'
utils.add_gzip_to_tar(string, path, tar)
# serialize weights
with utils.tmpfile() as tmppath:
torch.save(self.state_dict(), tmppath)
tar.add(tmppath, arcname='state_dict.pt')
# serialize current pie commit
if pie.__commit__ is not None:
string, path = pie.__commit__, 'pie-commit.zip'
utils.add_gzip_to_tar(string, path, tar)
# if passed, serialize settings
if settings is not None:
string, path = json.dumps(settings), 'settings.zip'
utils.add_gzip_to_tar(string, path, tar)
return fpath
def load(fpath):
"""
Load model from path
"""
import tarte.modules.models
with tarfile.open(utils.ensure_ext(fpath, 'tar'), 'r') as tar:
# load label encoder
le = MultiEncoder.load(
json.loads(utils.get_gzip_from_tar(tar, 'label_encoder.zip')))
# load model parameters
args, kwargs = json.loads(
utils.get_gzip_from_tar(tar, 'parameters.zip'))
# instantiate model
model_type = getattr(tarte.modules.models,
utils.get_gzip_from_tar(tar, 'class.zip'))
with utils.shutup():
model = model_type(le, *args, **kwargs)
# load settings
try:
settings = Settings(
json.loads(utils.get_gzip_from_tar(tar, 'settings.zip')))
model._settings = settings
except Exception:
logging.warn(
"Couldn't load settings for model {}!".format(fpath))
# load state_dict
with utils.tmpfile() as tmppath:
tar.extract('state_dict.pt', path=tmppath)
dictpath = os.path.join(tmppath, 'state_dict.pt')
model.load_state_dict(torch.load(dictpath, map_location='cpu'))
model.eval()
return model