Repository: facebookresearch/ParlAI

Issue:
MemNN dimension mismatch
I think that #1562 breaks MemNN.  I'm not sure I completely understand the problem, but it seems to be something along the lines of:

1. We build the dictionary given the training corpus
2. MemmNN adds some new tokens for its time features [[1](https://github.com/lematt1991/ParlAI/blob/master/parlai/agents/memnn/memnn.py#L92)].
3. We save the dictionary again, and it now has `memsize` extra tokens in the dictionary.
4. If you try to resume training or evaluate the model, when the model is instantiated it has an extra `memsize` features causing a dimension mismatch when loading the state dict.  

## Steps to repro

Train the model for at least `save_every_n_secs` .  Note that this will overwrite the original dictionary with an extra `memsize` tokens.

```
python examples/train_model.py --model memnn --task babi --save_every_n_secs 5 --position-encoding "True" -mf memnn/model --memsize 10
```

Train again resuming from the previous checkpoint (same command as above):

```
python examples/train_model.py --model memnn --task babi --save_every_n_secs 5 --position-encoding "True" -mf memnn/model --memsize 10
```

Results in:

```
RuntimeError: Error(s) in loading state_dict for MemNN:
        size mismatch for query_lt.weight: copying a param with shape torch.Size([211, 128]) from checkpoint, the shape in current model is torch.Size([243, 128]).
        size mismatch for in_memory_lt.weight: copying a param with shape torch.Size([211, 128]) from checkpoint, the shape in current model is torch.Size([243, 128]).
        size mismatch for out_memory_lt.weight: copying a param with shape torch.Size([211, 128]) from checkpoint, the shape in current model is torch.Size([243, 128]).
        size mismatch for answer_embedder.weight: copying a param with shape torch.Size([211, 128]) from checkpoint, theshape in current model is torch.Size([243, 128]).
(parlai) mattle@devfair0222:~/ParlAI$ python examples/train_model.py --model memnn --task babi --save_every_n_secs 5 --position-encoding "True" -mf memnn/model^C
(parlai) mattle@devfair0222:~/ParlAI$ python examples/train_model.py --model memnn --task babi --save_every_n_secs 5 --position-encoding "True" -mf memnn/model
```



Discussion:

Resolving patch:
diff --git a/parlai/agents/memnn/memnn.py b/parlai/agents/memnn/memnn.py
--- a/parlai/agents/memnn/memnn.py
+++ b/parlai/agents/memnn/memnn.py
@@ -63,22 +63,20 @@ def model_version():
         return 2
 
     def __init__(self, opt, shared=None):
-        # all instances may need some params
-        super().__init__(opt, shared)
-
         self.id = 'MemNN'
         self.memsize = opt['memsize']
         if self.memsize < 0:
             self.memsize = 0
         self.use_time_features = opt['time_features']
-
-        if not shared:
-            if opt['time_features']:
-                for i in range(self.memsize):
-                    self.dict[self._time_feature(i)] = 100000000 + i
+        super().__init__(opt, shared)
 
     def build_model(self):
         """Build MemNN model."""
+        if self.use_time_features:
+            # add time features to dictionary before building the model
+            for i in range(self.memsize):
+                self.dict[self._time_feature(i)] = 100000000 + i
+
         kwargs = opt_to_kwargs(self.opt)
         self.model = MemNN(len(self.dict), self.opt['embedding_size'],
                            padding_idx=self.NULL_IDX, **kwargs)
diff --git a/parlai/agents/memnn/modules.py b/parlai/agents/memnn/modules.py
--- a/parlai/agents/memnn/modules.py
+++ b/parlai/agents/memnn/modules.py
@@ -36,23 +36,10 @@ def __init__(
         # prepare features
         self.hops = hops
 
-        # time features: we learn an embedding for each memory slot
-        self.extra_features = 0
-        if time_features:
-            self.extra_features += memsize
-            self.time_features = torch.LongTensor(
-                range(num_features, num_features + memsize))
-
         def embedding(use_extra_feats=True):
-            if use_extra_feats:
-                return Embed(num_features + self.extra_features,
-                             embedding_size,
-                             position_encoding=position_encoding,
-                             padding_idx=padding_idx)
-            else:
-                return Embed(num_features, embedding_size,
-                             position_encoding=position_encoding,
-                             padding_idx=padding_idx)
+            return Embed(num_features, embedding_size,
+                         position_encoding=position_encoding,
+                         padding_idx=padding_idx)
 
         # TODO: add token dropout?
         # TODO: add dropout


Test patch: