| 29 | |
| 30 | |
| 31 | class KestrelSolver(ClsSolver): |
| 32 | |
| 33 | def __init__(self, config_file, recover=''): |
| 34 | self.config_file = config_file |
| 35 | self.recover = recover |
| 36 | self.prototype_info = EasyDict() |
| 37 | self.config = parse_config(config_file) |
| 38 | self.setup_env() |
| 39 | self.build_model() |
| 40 | # 'recover' only for convert |
| 41 | if self.recover: |
| 42 | self.logger.info(f"Recover exist! Again Recovering from {self.recover}") |
| 43 | recover_state = torch.load(self.recover, 'cpu') |
| 44 | load_state_model(self.model, recover_state['model']) |
| 45 | if self.config.to_kestrel.get('add_softmax'): |
| 46 | self.model = Wrapper(self.model) |
| 47 | |
| 48 | def to_caffe(self, save_prefix='model', input_size=None): |
| 49 | try: |
| 50 | from spring.nart.tools import pytorch |
| 51 | except ImportError: |
| 52 | print('Install Spring NART first!') |
| 53 | |
| 54 | with pytorch.convert_mode(): |
| 55 | pytorch.convert(self.model.float(), |
| 56 | [(3, self.config.data.input_size, |
| 57 | self.config.data.input_size)], |
| 58 | filename=save_prefix, |
| 59 | input_names=['data'], |
| 60 | output_names=['out']) |
| 61 | |
| 62 | def to_nnie(self, nnie_cfg, config, prototxt, caffemodel, model_name): |
| 63 | nnie_cfg_path = generate_nnie_config(nnie_cfg, config) |
| 64 | nnie_cmd = 'python -m spring.nart.switch -c {} -t nnie {} {}'.format( |
| 65 | nnie_cfg_path, prototxt, caffemodel) |
| 66 | |
| 67 | os.system(nnie_cmd) |
| 68 | assert os.path.exists("parameters.json") |
| 69 | with open("parameters.json", "r") as f: |
| 70 | params = json.load(f) |
| 71 | params["model_files"]["net"]["net"] = "engine.bin" |
| 72 | params["model_files"]["net"]["backend"] = "kestrel_nart" |
| 73 | with open("parameters.json", "w") as f: |
| 74 | json.dump(params, f, indent=2) |
| 75 | tar_cmd = 'tar cvf {} engine.bin engine.bin.json meta.json meta.conf parameters.json category_param.json'.\ |
| 76 | format(model_name + "_nnie.tar") |
| 77 | os.system(tar_cmd) |
| 78 | self.logger.info(f"generate {model_name + '_nnie.tar'} done!") |
| 79 | |
| 80 | def refactor_config(self): |
| 81 | '''Prepare configuration for kestrel classifier model. For details: |
| 82 | https://confluence.sensetime.com/display/VIBT/nart.tools.kestrel.classifier |
| 83 | ''' |
| 84 | kestrel_config = EasyDict() |
| 85 | kestrel_config['pixel_means'] = self.config.to_kestrel.get('pixel_means', [123.675, 116.28, 103.53]) |
| 86 | kestrel_config['pixel_stds'] = self.config.to_kestrel.get('pixel_stds', [58.395, 57.12, 57.375]) |
| 87 | |
| 88 | kestrel_config['is_rgb'] = self.config.to_kestrel.get('is_rgb', True) |