def _torch_snapshot_object(trainer, target, filename, savefun):
from chainer.serializers import DictionarySerializer
# make snapshot_dict dictionary
s = DictionarySerializer()
s.save(trainer)
if hasattr(trainer.updater.model, "model"):
# (for TTS)
if hasattr(trainer.updater.model.model, "module"):
model_state_dict = trainer.updater.model.model.module.state_dict()
else:
model_state_dict = trainer.updater.model.model.state_dict()
else:
# (for ASR)
if hasattr(trainer.updater.model, "module"):
model_state_dict = trainer.updater.model.module.state_dict()
else:
model_state_dict = trainer.updater.model.state_dict()
snapshot_dict = {
"trainer": s.target,
"model": model_state_dict,
"optimizer": trainer.updater.get_optimizer("main").state_dict(),
}
# save snapshot dictionary
fn = filename.format(trainer)
prefix = "tmp" + fn
tmpdir = tempfile.mkdtemp(prefix=prefix, dir=trainer.out)
tmppath = os.path.join(tmpdir, fn)
try:
savefun(snapshot_dict, tmppath)
shutil.move(tmppath, os.path.join(trainer.out, fn))
finally:
shutil.rmtree(tmpdir)
def test_serialize():
# Sample dataset
dataset = get_dataset()
# Iterator
it = LtrIterator(dataset, repeat=True, shuffle=False)
# Set up serializer
serializer = DictionarySerializer()
# Serialize
it.serialize(serializer)
print(serializer.target['epoch'])
# Perform one epoch iteration
it.next()
it.next()
it.next()
# Before serializing, our epoch variable should now be different
assert_not_equal(serializer.target['epoch'], it.epoch)
# After serializing it should be equal again
it.serialize(serializer)
assert_equal(serializer.target['epoch'], it.epoch)
def save_non_embed_npz(file, obj, compression=True):
"""
Since we do not update the original embedding matrix,
it is unnecessary (and waste of storage) to include it in model file.
--> Exclude embed_mat from .npz file
"""
if isinstance(file, six.string_types):
with open(file, 'wb') as f:
save_non_embed_npz(f, obj, compression)
return
s = DictionarySerializer()
s.save(obj)
s.target = {k: v for k, v in s.target.items() if 'embed_mat' not in k}
if compression:
np.savez_compressed(file, **s.target)
else:
np.savez(file, **s.target)