Load a model (identified by the config used for construction) and return it
(self, logdir, step)
| 36 | self.model_preproc.load() |
| 37 | |
| 38 | def load_model(self, logdir, step): |
| 39 | '''Load a model (identified by the config used for construction) and return it''' |
| 40 | # 1. Construct model |
| 41 | model = registry.construct('model', self.config['model'], preproc=self.model_preproc, device=self.device) |
| 42 | model.to(self.device) |
| 43 | model.eval() |
| 44 | model.visualize_flag = False |
| 45 | |
| 46 | # 2. Restore its parameters |
| 47 | saver = saver_mod.Saver({"model": model}) |
| 48 | last_step = saver.restore(logdir, step=step, map_location=self.device, item_keys=["model"]) |
| 49 | |
| 50 | if not last_step: |
| 51 | raise Exception('Attempting to infer on untrained model') |
| 52 | return model |
| 53 | |
| 54 | def infer(self, model, output_path, args): |
| 55 | output = open(output_path, 'w') |