Reset parameters if hidden size is changed

hm-ysjiang · Jun 7, 2023 · c41db70 · c41db70
1 parent b7546ea
commit c41db70
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/train-selfsupervised.py b/train-selfsupervised.py
@@ -116,7 +116,13 @@ def train(args):
         if args.reset_context:
             _weight = OrderedDict()
             for key, val in checkpoint.items():
-                if '.cnet.' not in key and '.update_block.gru.' not in key:
+                if args.context != 128 and \
+                    ('.cnet.' in key or '.update_block.gru.' in key):
+                    pass
+                elif args.hidden != 128 and \
+                    ('.update_block.gru.' in key or '.update_block.flow_head.' in key):
+                    pass
+                else:
                     _weight[key] = val
             weight = _weight
         model.load_state_dict(weight, strict=(not args.allow_nonstrict))

diff --git a/train-supervised.py b/train-supervised.py
@@ -110,7 +110,13 @@ def train(args):
         if args.reset_context:
             _weight = OrderedDict()
             for key, val in checkpoint.items():
-                if '.cnet.' not in key and '.update_block.gru.' not in key:
+                if args.context != 128 and \
+                    ('.cnet.' in key or '.update_block.gru.' in key):
+                    pass
+                elif args.hidden != 128 and \
+                    ('.update_block.gru.' in key or '.update_block.flow_head.' in key):
+                    pass
+                else:
                     _weight[key] = val
             weight = _weight
         model.load_state_dict(weight, strict=(not args.allow_nonstrict))