Rework how weights and options are stored in the .npz file

ozancaglayan · ozancaglayan · commit 0721f34924d2 · 2017-10-16T15:12:14.000+02:00
Previously the storing was not efficient causing %30 bigger .npz files.
This is now fixed. Always use get_param_dict() and get_model_options()
methods from sysutils to get the options and tparams of your model
files.
diff --git a/bin/nmt-extract b/bin/nmt-extract
@@ -7,7 +7,7 @@ import argparse
 
 import numpy as np
 
-from nmtpy.nmtutils import get_param_dict
+from nmtpy.sysutils import get_param_dict
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(prog='nmt-extract')
diff --git a/bin/nmt-rescore b/bin/nmt-rescore
@@ -87,7 +87,7 @@ if __name__ == "__main__":
     log.info("%d models given for rescoring" % len(args.models))
 
     # Load model options from first model
-    model_options = dict(np.load(args.models[0])['opts'].tolist())
+    model_options = get_model_options(args.models[0])
 
     # Import the module
     Model = importlib.import_module("nmtpy.models.%s" % model_options['model_type']).Model
@@ -138,7 +138,7 @@ if __name__ == "__main__":
 
         # Load model weights for anything except first one (hacky)
         if idx > 0:
-            model.update_shared_variables(np.load(modelfile)['tparams'].tolist())
+            model.update_shared_variables(get_param_dict(modelfile))
 
         for i, data in enumerate(iterator):
             norm = data['y_mask'].sum(0)
diff --git a/bin/nmt-test-lm b/bin/nmt-test-lm
@@ -59,7 +59,7 @@ class Tester(object):
         self.processes = [None] * self.n_jobs
 
     def set_model_options(self):
-        model_options = dict(np.load(self.model_file)['opts'].tolist())
+        model_options = get_model_options(self.model_file)
 
         # Import the module
         self.__class = importlib.import_module("nmtpy.models.%s" % model_options['model_type']).Model
diff --git a/bin/nmt-train b/bin/nmt-train
@@ -13,7 +13,6 @@ import importlib
 from nmtpy.config import Config
 from nmtpy.logger import Logger
 from nmtpy.sysutils import *
-from nmtpy.nmtutils import get_param_dict
 from nmtpy.mainloop import MainLoop
 
 # Ensure cleaning up temp files and processes
diff --git a/bin/nmt-translate b/bin/nmt-translate
@@ -120,7 +120,7 @@ class Translator(object):
             data = np.load(mfile)
 
             # Fetch options
-            model_options = get_model_options(data['opts'])
+            model_options = get_model_options(data)
 
             # Import the module
             self.__class = importlib.import_module("nmtpy.models.%s" % model_options['model_type']).Model
@@ -129,7 +129,7 @@ class Translator(object):
             model = self.__class(seed=self.seed, logger=None, **model_options)
 
             # Load weights from the already loaded file
-            model.load(data['tparams'].tolist())
+            model.load(data)
 
             model.set_dropout(False)
             model.build_sampler()
diff --git a/bin/nmt-translate-factors b/bin/nmt-translate-factors
@@ -116,7 +116,7 @@ class Translator(object):
             log.info('Initializing model %s' % os.path.basename(mfile))
 
             # Fetch options
-            model_options = dict(np.load(mfile)['opts'].tolist())
+            model_options = get_model_options(mfile)
 
             # Import the module
             self.__class = importlib.import_module("nmtpy.models.%s" % model_options['model_type']).Model
diff --git a/bin/nmt-translate-server b/bin/nmt-translate-server
@@ -120,16 +120,14 @@ class Translator(object):
             # Load model file
             data = np.load(mfile)
 
-            #model_options = dict(np.load(mfile)['opts'].tolist())
-            model_options = get_model_options(data['opts'])
+            model_options = get_model_options(data)
 
             # Import the module
             self.__class = importlib.import_module("nmtpy.models.%s" % model_options['model_type']).Model
 
             # Create the model
             model = self.__class(seed=self.seed, logger=None, **model_options)
-            #model.load(mfile)
-            model.load(data['tparams'].tolist())
+            model.load(data)
             model.set_dropout(False)
             model.build_sampler()
 
diff --git a/nmtpy/models/basemodel.py b/nmtpy/models/basemodel.py
@@ -8,8 +8,8 @@
 from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
 
 import numpy as np
-from ..nmtutils import unzip, get_param_dict
-from ..sysutils import readable_size, get_temp_file, get_valid_evaluation
+from ..nmtutils import unzip
+from ..sysutils import readable_size, get_temp_file, get_param_dict, get_valid_evaluation
 from ..defaults import INT, FLOAT
 from ..optimizers import get_optimizer
 
@@ -83,18 +83,19 @@ def get_nb_params(self):
 
     def save(self, fname):
         """Save model parameters as .npz."""
+        kwargs = OrderedDict()
+        kwargs['opts'] = self._options
         if self.tparams is not None:
-            np.savez(fname, tparams=unzip(self.tparams), opts=self._options)
-        else:
-            np.savez(fname, opts=self._options)
+            kwargs.update(unzip(self.tparams))
+
+        # Save each param as a separate argument into npz
+        np.savez(fname, **kwargs)
 
     def load(self, params):
         """Restore .npz checkpoint file into model."""
         self.tparams = OrderedDict()
 
-        if isinstance(params, str):
-            # Filename, load from it
-            params = get_param_dict(params)
+        params = get_param_dict(params)
 
         for k,v in params.items():
             self.tparams[k] = theano.shared(v.astype(FLOAT), name=k)
diff --git a/nmtpy/nmtutils.py b/nmtpy/nmtutils.py
@@ -53,10 +53,6 @@ def unzip(zipped):
 def pp(prefix, name):
     return '%s_%s' % (prefix, name)
 
-def get_param_dict(path):
-    """Fetch parameter dictionary from .npz file."""
-    return np.load(path)['tparams'].tolist()
-
 # orthogonal initialization for weights
 # Saxe, Andrew M., James L. McClelland, and Surya Ganguli.
 # "Exact solutions to the nonlinear dynamics of learning in deep
diff --git a/nmtpy/sysutils.py b/nmtpy/sysutils.py
@@ -8,6 +8,10 @@
 import tempfile
 import subprocess
 
+from collections import OrderedDict
+
+import numpy as np
+
 from . import cleanup
 
 def print_summary(train_args, model_args, print_func=None):
@@ -293,13 +297,37 @@ def get_next_runid(save_path, exp_id):
 
     return i
 
-def get_model_options(optdict):
+def get_model_options(handle):
     """Get model options by maintaining backward compatibility."""
-    optdict = optdict.tolist()
+    if isinstance(handle, np.lib.npyio.NpzFile):
+        optdict = handle['opts'].tolist()
+    elif isinstance(handle, str):
+        optdict = np.load(path)['opts'].tolist()
+    elif isinstance(handle, np.ndarray):
+        optdict = handle.tolist()
+    elif isinstance(handle, dict):
+        optdict = handle
 
     # tied_trg_emb, i.e. renamed to tied_emb
     if 'tied_trg_emb' in optdict and 'tied_emb' not in optdict:
         optdict['tied_emb'] = '2way'
         del optdict['tied_trg_emb']
 
     return optdict
+
+def get_param_dict(handle):
+    """Fetch parameter dictionary from .npz file."""
+    if isinstance(handle, np.lib.npyio.NpzFile):
+        npz = handle
+    elif isinstance(handle, str):
+        npz = np.load(path)
+    elif isinstance(handle, dict):
+        return handle
+
+    if 'tparams' in npz.files:
+        # Old version of storing things
+        return npz['tparams'].tolist()
+    else:
+        params = OrderedDict(npz.iteritems())
+        del params['opts']
+        return params
diff --git a/scripts/modify-npz b/scripts/modify-npz
@@ -6,6 +6,9 @@ import sys
 import numpy as np
 
 from ast import literal_eval
+from collections import OrderedDict
+
+from nmtpy.sysutils import *
 
 os.environ['THEANO_FLAGS'] = 'device=cpu'
 
@@ -31,8 +34,10 @@ if __name__ == '__main__':
     for npzf in files:
         # Load the file
         npz = np.load(npzf)
+
         # Get the dict
-        opts = npz['opts'].tolist()
+        opts = get_model_options(npz)
+        params = get_param_dict(npz)
 
         newfilename = npzf
 
@@ -46,4 +51,7 @@ if __name__ == '__main__':
                 newfilename = '%s-%s' % (value, rest)
 
         print('Writing %s' % newfilename)
-        np.savez(newfilename, tparams=npz['tparams'], opts=opts)
+
+        params['opts'] = opts
+
+        np.savez(newfilename, **params)
