Merge pull request #101 from lucasb-eyer/safe-setstate

A whole bunch of sanity-checks in parameter loading.

Author: lucasb-eyer

DeepFried2/Container.py

@@ -64,6 +64,9 @@ def __getstate__(self):
         return [m.__getstate__() for m in self.modules]
 
     def __setstate__(self, state):
+        if len(self.modules) != len(state):
+            raise ValueError("{} wants to load params for {} modules but received params for {} modules".format(df.utils.typename(self), len(self.modules), len(state)))
+
         for m, s in zip(self.modules, state):
             m.__setstate__(s)
 

DeepFried2/Module.py

@@ -189,5 +189,12 @@ def __getstate__(self):
         return [p.get_value() for p in self.parameters()]
 
     def __setstate__(self, state):
-        for p, s in zip(self.parameters(), state):
+        params = self.parameters()
+        if len(params) != len(state):
+            raise ValueError("{} wants to load {} params but received {} params".format(df.utils.typename(self), len(params), len(state)))
+
+        for p, s in zip(params, state):
+            if p.get_value().shape != s.shape:
+                raise ValueError("{} got invalid shape when loading param {}: expecting {} but loading {}".format(df.utils.typename(self), p.param.name, p.get_value().shape, s.shape))
+
             p.set_value(s)

DeepFried2/layers/BatchNormalization.py

@@ -107,7 +107,6 @@ def __getstate__(self):
         return [buf.get_value() for buf in (self.buf_mean, self.buf_var, self.buf_count)] + regular
 
     def __setstate__(self, state):
-        istate = iter(state)
-        for buf, val in zip((self.buf_mean, self.buf_var, self.buf_count), istate):
+        for buf, val in zip((self.buf_mean, self.buf_var, self.buf_count), state):
             buf.set_value(val)
-        df.Module.__setstate__(self, istate)
+        df.Module.__setstate__(self, state[3:])