| 21 | from seq2struct.models.spider import spider_beam_search |
| 22 | |
| 23 | class Inferer: |
| 24 | def __init__(self, config): |
| 25 | self.config = config |
| 26 | if torch.cuda.is_available(): |
| 27 | self.device = torch.device('cuda') |
| 28 | else: |
| 29 | self.device = torch.device('cpu') |
| 30 | torch.set_num_threads(1) |
| 31 | |
| 32 | # 0. Construct preprocessors |
| 33 | self.model_preproc = registry.instantiate( |
| 34 | registry.lookup('model', config['model']).Preproc, |
| 35 | config['model']) |
| 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') |
| 56 | |
| 57 | with torch.no_grad(): |
| 58 | if args.mode == 'infer': |
| 59 | orig_data = registry.construct('dataset', self.config['data'][args.section]) |
| 60 | preproc_data = self.model_preproc.dataset(args.section) |
| 61 | if args.limit: |
| 62 | sliced_orig_data = itertools.islice(orig_data, args.limit) |
| 63 | sliced_preproc_data = itertools.islice(preproc_data, args.limit) |
| 64 | else: |
| 65 | sliced_orig_data = orig_data |
| 66 | sliced_preproc_data = preproc_data |
| 67 | assert len(orig_data) == len(preproc_data) |
| 68 | self._inner_infer(model, args.beam_size, args.output_history, sliced_orig_data, sliced_preproc_data, output, args.use_heuristic) |
| 69 | elif args.mode == 'debug': |
| 70 | data = self.model_preproc.dataset(args.section) |
| 71 | if args.limit: |
| 72 | sliced_data = itertools.islice(data, args.limit) |
| 73 | else: |
| 74 | sliced_data = data |
| 75 | self._debug(model, sliced_data, output) |
| 76 | elif args.mode == 'visualize_attention': |
| 77 | model.visualize_flag = True |
| 78 | model.decoder.visualize_flag = True |
| 79 | data = registry.construct('dataset', self.config['data'][args.section]) |
| 80 | if args.limit: |