Fix a bug

zsLin177 · zsLin177 · commit 62f660ceddcc · 2021-06-16T13:25:58.000+08:00
diff --git a/supar/__init__.py b/supar/__init__.py
@@ -65,3 +65,6 @@
 }
 MODEL = {n: f"{SRC['github']}/v1.1.0/{m}.zip" for n, m in NAME.items()}
 CONFIG = {n: f"{SRC['github']}/v1.1.0/{m}.ini" for n, m in NAME.items()}
+
+MODEL['biaffine-sdp-en'] = f"{SRC['hlt']}/v1.1.1/{NAME['biaffine-sdp-en']}.zip"
+MODEL['biaffine-sdp-zh'] = f"{SRC['hlt']}/v1.1.1/{NAME['biaffine-sdp-zh']}.zip"
diff --git a/supar/cmds/cmd.py b/supar/cmds/cmd.py
@@ -25,11 +25,11 @@ def parse(parser):
     logger.info('\n' + str(args))
 
     if args.mode == 'train':
-        parser = Parser.load(args.path) if args.checkpoint else Parser.build(**args)
+        parser = Parser.load(**args) if args.checkpoint else Parser.build(**args)
         parser.train(**args)
     elif args.mode == 'evaluate':
-        parser = Parser.load(args.path)
+        parser = Parser.load(**args)
         parser.evaluate(**args)
     elif args.mode == 'predict':
-        parser = Parser.load(args.path)
+        parser = Parser.load(**args)
         parser.predict(**args)
diff --git a/supar/parsers/parser.py b/supar/parsers/parser.py
@@ -62,11 +62,11 @@ def train(self, train, dev, test, buckets=32, batch_size=5000, update_steps=1,
 
         self.epoch, self.best_e, self.patience, self.best_metric, self.elapsed = 1, 1, patience, Metric(), timedelta()
         if self.args.checkpoint:
-            self.optimizer.load_state_dict(self.args.pop('optimizer_state_dict'))
-            self.scheduler.load_state_dict(self.args.pop('scheduler_state_dict'))
-            for k, v in args.pop('state_args').items():
+            self.optimizer.load_state_dict(self.checkpoint_state_dict.pop('optimizer_state_dict'))
+            self.scheduler.load_state_dict(self.checkpoint_state_dict.pop('scheduler_state_dict'))
+            set_rng_state(self.checkpoint_state_dict.pop('rng_state'))
+            for k, v in self.checkpoint_state_dict.items():
                 setattr(self, k, v)
-            set_rng_state(args.pop('rng_state'))
             train.loader.batch_sampler.epoch = self.epoch
 
         for epoch in range(self.epoch, args.epochs + 1):
@@ -202,7 +202,9 @@ def load(cls, path, reload=False, src=None, **kwargs):
         model.load_state_dict(state['state_dict'], False)
         model.to(args.device)
         transform = state['transform']
-        return cls(args, model, transform)
+        parser = cls(args, model, transform)
+        parser.checkpoint_state_dict = state['checkpoint_state_dict'] if args.checkpoint else None
+        return parser
 
     def save(self, path):
         model = self.model
@@ -223,15 +225,16 @@ def save_checkpoint(self, path):
         if hasattr(model, 'module'):
             model = self.model.module
         args = model.args
-        args.state_args = {k: getattr(self, k) for k in ['epoch', 'best_e', 'patience', 'best_metric', 'elapsed']}
-        args.optimizer_state_dict = self.optimizer.state_dict()
-        args.scheduler_state_dict = self.scheduler.state_dict()
-        args.rng_state = get_rng_state()
+        checkpoint_state_dict = {k: getattr(self, k) for k in ['epoch', 'best_e', 'patience', 'best_metric', 'elapsed']}
+        checkpoint_state_dict.update({'optimizer_state_dict': self.optimizer.state_dict(),
+                                      'scheduler_state_dict': self.scheduler.state_dict(),
+                                      'rng_state': get_rng_state()})
         state_dict = {k: v.cpu() for k, v in model.state_dict().items()}
         pretrained = state_dict.pop('pretrained.weight', None)
         state = {'name': self.NAME,
                  'args': args,
                  'state_dict': state_dict,
                  'pretrained': pretrained,
+                 'checkpoint_state_dict': checkpoint_state_dict,
                  'transform': self.transform}
         torch.save(state, path, pickle_module=dill)

Original file line number	Diff line number	Diff line change
`@@ -65,3 +65,6 @@`
`65`	`65`	`}`
`66`	`66`	`MODEL = {n: f"{SRC['github']}/v1.1.0/{m}.zip" for n, m in NAME.items()}`
`67`	`67`	`CONFIG = {n: f"{SRC['github']}/v1.1.0/{m}.ini" for n, m in NAME.items()}`
	`68`	`+`
	`69`	`+MODEL['biaffine-sdp-en'] = f"{SRC['hlt']}/v1.1.1/{NAME['biaffine-sdp-en']}.zip"`
	`70`	`+MODEL['biaffine-sdp-zh'] = f"{SRC['hlt']}/v1.1.1/{NAME['biaffine-sdp-zh']}.zip"`