import torch import skyhook.pytorch.model as model import skyhook.pytorch.util as util in_dataset_id = int(sys.argv[1]) device = torch.device('cuda:0') #device = torch.device('cpu') model_path = 'data/items/{}/model.pt'.format(in_dataset_id) save_dict = torch.load(model_path) example_inputs = save_dict['example_inputs'] util.inputs_to_device(example_inputs, device) net = model.Net(save_dict['arch'], save_dict['comps'], example_inputs, save_dict['example_metadatas'], infer=True, device=device) net.to(device) net.load_state_dict(save_dict['model']) net.eval() stdin = sys.stdin.detach() while True: header = stdin.read(8) if not header: break left_count, right_count = struct.unpack('>II', header) buf = stdin.read(left_count * 64 * 64 * 3) left_arr = numpy.frombuffer(buf, dtype='uint8').reshape(
comp_params = json.loads(arch['Components'][0].get('Params', '{}')) comp_params['mode'] = mode arch['Components'][0]['Params'] = json.dumps(comp_params) # example inputs im_data = numpy.zeros((416, 416, 3), dtype='uint8') example_inputs = [ util.collate('image', [util.prepare_input('image', im_data, {}, {})]), util.collate( 'detection', [util.prepare_input('detection', [], {'CanvasDims': [416, 416]}, {})]), ] util.inputs_to_device(example_inputs, device) # example metadata with open(os.path.join(yolo_path, 'data', 'coco.yaml'), 'r') as f: d = yaml.load(f, Loader=yaml.FullLoader) categories = d['names'] example_metadatas = [{}, {'Categories': categories}] net = model.Net(arch, comps, example_inputs, example_metadatas, device=device) sys.path.append(yolo_path) orig_dict = torch.load(in_fname)['model'].state_dict() state_dict = {} for k, v in orig_dict.items(): state_dict['mlist.0.model.' + k] = v net.load_state_dict(state_dict) torch.save(net.get_save_dict(), out_fname)
} for comp_idx, comp_spec in enumerate(arch['Components']): comp_params = {} if comp_spec['Params']: comp_params = json.loads(comp_spec['Params']) if overwrite_comp_params.get(comp_idx, None): comp_params.update(json.loads(overwrite_comp_params[comp_idx])) comp_spec['Params'] = json.dumps(comp_params) example_inputs = save_dict['example_inputs'] util.inputs_to_device(example_inputs, device) net = model.Net(arch, save_dict['comps'], example_inputs, save_dict['example_metadatas'], output_datasets=params['OutputDatasets'], infer=True, device=device) net.to(device) net.load_state_dict(save_dict['model']) net.eval() input_options = {} for spec in params['InputOptions']: input_options[spec['Idx']] = json.loads(spec['Value']) meta = None