def main(): args = parse_args() out_name = args.out cfg_path = args.config cfg = Config.fromfile(cfg_path) deploy_cfg = cfg['deploy'] common_cfg = cfg.get('common') runner = InferenceRunner(deploy_cfg, common_cfg) assert runner.use_gpu, 'Please use valid gpu to export model.' runner.load_checkpoint(args.checkpoint) image = cv2.imread(args.image) aug = runner.transform(image=image, label='') image, label = aug['image'], aug['label'] image = image.unsqueeze(0).cuda() dummy_input = (image, runner.converter.test_encode([''])) model = runner.model.cuda().eval() need_text = runner.need_text if not need_text: dummy_input = dummy_input[0] if args.onnx: runner.logger.info('Convert to onnx model') torch2onnx(model, dummy_input, out_name) else: max_batch_size = args.max_batch_size max_workspace_size = args.max_workspace_size fp16_mode = args.fp16 int8_mode = args.int8 int8_calibrator = None if int8_mode: runner.logger.info('Convert to trt engine with int8') if args.calibration_images: runner.logger.info('Use calibration with mode {} and data {}' .format(args.calibration_mode, args.calibration_images)) dataset = CalibDataset(args.calibration_images, runner.converter, runner.transform, need_text) int8_calibrator = CALIBRATORS[args.calibration_mode]( dataset=dataset) else: runner.logger.info('Use default calibration mode and data') elif fp16_mode: runner.logger.info('Convert to trt engine with fp16') else: runner.logger.info('Convert to trt engine with fp32') trt_model = torch2trt( model, dummy_input, max_batch_size=max_batch_size, max_workspace_size=max_workspace_size, fp16_mode=fp16_mode, int8_mode=int8_mode, int8_calibrator=int8_calibrator) save(trt_model, out_name) runner.logger.info( 'Convert successfully, save model to {}'.format(out_name))
def get_model(file_config, weights): cfg = Config.fromfile(file_config) deploy_cfg = cfg['deploy'] model = build_from_cfg(deploy_cfg['model'], MODELS) checkpoint = torch.load(weights, map_location='cpu') model.load_state_dict(checkpoint['state_dict']) return model
def __init__(self, config, weights) -> None: super().__init__() self.config = config self.weights = weights cfg = Config.fromfile(self.config) deploy_cfg = cfg['deploy'] common_cfg = cfg.get('common') cfg['batch_max_length'] = 40 runner = InferenceRunner(deploy_cfg, common_cfg) runner.load_checkpoint(self.weights) self.runner = runner
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) deploy_cfg = cfg['deploy'] common_cfg = cfg.get('common') if torch.cuda.is_available(): device = torch.cuda.current_device() deploy_cfg['gpu_id'] = str(device) else: raise AssertionError('Please use gpu for benchmark.') runner = InferenceRunner(deploy_cfg, common_cfg) runner.load_checkpoint(args.checkpoint) C, H, W = [int(_.strip()) for _ in args.dummy_input_shape.split(',')] dummy_image = np.random.random_integers(0, 255, (H, W, C)).astype(np.uint8) aug = runner.transform(image=dummy_image, label='') image, label = aug['image'], aug['label'] image = image.unsqueeze(0).cuda() dummy_input = (image, runner.converter.test_encode([''])) model = runner.model.cuda().eval() need_text = runner.need_text if not need_text: dummy_input = dummy_input[0] if args.dynamic_shape: print(f'Convert to Onnx with dynamic input shape and opset version' f'{args.opset_version}') else: print(f'Convert to Onnx with constant input shape' f' {args.dummy_input_shape} and opset version ' f'{args.opset_version}') torch2onnx( model, dummy_input, args.out, verbose=args.verbose, dynamic_shape=args.dynamic_shape, opset_version=args.opset_version, do_constant_folding=args.do_constant_folding, ) runner.logger.info( f'Convert successfully, saved onnx file: {os.path.abspath(args.out)}')
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) test_cfg = cfg['test'] deploy_cfg = cfg['deploy'] common_cfg = cfg['common'] runner = TestRunner(test_cfg, deploy_cfg, common_cfg) assert runner.use_gpu, 'Please use gpu for benchmark.' runner.load_checkpoint(args.checkpoint) # image = Image.open(args.image) image = cv2.imread(args.image) aug= runner.transform(image=image,label='1') image, dummy_label = aug['image'], aug['label'] image = image.unsqueeze(0) input_len = runner.converter.test_encode(1)[0] model = runner.model need_text = runner.need_text if need_text: shape = tuple(image.shape), tuple(input_len.shape) else: shape = tuple(image.shape) dtypes = args.dtypes iters = args.iters int8_calibrator = None if args.calibration_images: calib_dataset = CalibDataset(args.calibration_images, runner.converter, runner.transform, need_text) int8_calibrator = [CALIBRATORS[mode](dataset=calib_dataset) for mode in args.calibration_modes] # assert isinstance(runner.test_dataloader['all'], tud.DataLoader), \ # "Only suppor single dataloader in training phase. " \ # "Check the type of dataset please. " \ # "If the type of dataset is list, then the type of build datalodaer will be dict." \ # "If the type of dataset is torch.utils.data.Dataset, " \ # "the type of build dataloader will be torch.utils.data.Dataloader. " \ # "If you wanna combine different dataset, consider using ConcatDataset in your config file please." # dataset = runner.test_dataloader['all'].dataset # dataset = MetricDataset(dataset, runner.converter, need_text) metric = Metric(runner.metric, runner.converter) benchmark(model, shape, dtypes=dtypes, iters=iters, int8_calibrator=int8_calibrator, metric=metric)
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) test_cfg = cfg['test'] deploy_cfg = cfg['deploy'] common_cfg = cfg['common'] runner = TestRunner(test_cfg, deploy_cfg, common_cfg) assert runner.use_gpu, 'Please use gpu for benchmark.' runner.load_checkpoint(args.checkpoint) image = Image.open(args.image) image, dummy_label = runner.transform(image, '') image = image.unsqueeze(0) input_len = runner.converter.test_encode(1)[0] model = runner.model need_text = runner.need_text if need_text: shape = tuple(image.shape), tuple(input_len.shape) else: shape = tuple(image.shape) dtypes = args.dtypes iters = args.iters int8_calibrator = None if args.calibration_images: calib_dataset = CalibDataset(args.calibration_images, runner.converter, runner.transform, need_text) int8_calibrator = [ CALIBRATORS[mode](dataset=calib_dataset) for mode in args.calibration_modes ] dataset = runner.test_dataloader.dataset dataset = MetricDataset(dataset, runner.converter, need_text) metric = Metric(runner.metric, runner.converter) benchmark(model, shape, dtypes=dtypes, iters=iters, int8_calibrator=int8_calibrator, dataset=dataset, metric=metric)
def load_model(): cfg_path = os.path.join(os.path.dirname(__file__), "vedastr/configs/small_satrn.py") cfg = Config.fromfile(cfg_path) deploy_cfg = cfg['deploy'] common_cfg = cfg.get('common') checkpoint = os.path.join(home, *["OCR", "vedastr", "small_satrn.pth"]) if not os.path.exists(checkpoint): if not os.path.exists(os.path.dirname(checkpoint)): os.makedirs(os.path.dirname(checkpoint)) download_drive_file(file_id="1bcKtEcYGIOehgPfGi_TqPkvrm6rjOUKR", output=checkpoint) runner = InferenceRunner(deploy_cfg, common_cfg) runner.load_checkpoint(checkpoint) # print("model loaded..") return runner
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) _, fullname = os.path.split(cfg_path) fname, ext = os.path.splitext(fullname) root_workdir = cfg.pop('root_workdir') workdir = os.path.join(root_workdir, fname) os.makedirs(workdir, exist_ok=True) train_cfg = cfg['train'] deploy_cfg = cfg['deploy'] common_cfg = cfg['common'] common_cfg['workdir'] = workdir runner = TrainRunner(train_cfg, deploy_cfg, common_cfg) runner()
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) deploy_cfg = cfg['deploy'] common_cfg = cfg.get('common') runner = DeployRunner(deploy_cfg, common_cfg) runner.load_checkpoint(args.checkpoint) if os.path.isfile(args.image): images = [args.image] else: images = [ os.path.join(args.image, name) for name in os.listdir(args.image) ] for img in images: image = Image.open(img) pred_str, probs = runner(image) runner.logger.info('predict string: {} \t of {}'.format(pred_str, img))
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) _, fullname = os.path.split(cfg_path) fname, ext = os.path.splitext(fullname) root_workdir = cfg.pop('root_workdir') workdir = os.path.join(root_workdir, fname) os.makedirs(workdir, exist_ok=True) test_cfg = cfg['test'] deploy_cfg = cfg['deploy'] common_cfg = cfg['common'] common_cfg['workdir'] = workdir deploy_cfg['gpu_id'] = args.gpus.replace(" ", "") runner = TestRunner(test_cfg, deploy_cfg, common_cfg) runner.load_checkpoint(args.checkpoint) runner()
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) deploy_cfg = cfg['deploy'] common_cfg = cfg.get('common') deploy_cfg['gpu_id'] = args.gpus.replace(" ", "") runner = InferenceRunner(deploy_cfg, common_cfg) runner.load_checkpoint(args.checkpoint) if os.path.isfile(args.image): images = [args.image] else: images = [ os.path.join(args.image, name) for name in os.listdir(args.image) ] for img in images: image = cv2.imread(img) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) pred_str, probs = runner(image) runner.logger.info('predict string: {} \t of {}'.format(pred_str, img))
def main(): args = parse_args() cfg_path = args.config cfg = Config.fromfile(cfg_path) test_cfg = cfg['test'] deploy_cfg = cfg['deploy'] common_cfg = cfg['common'] if torch.cuda.is_available(): device = torch.cuda.current_device() deploy_cfg['gpu_id'] = str(device) else: raise AssertionError('Please use gpu for benchmark.') runner = TestRunner(test_cfg, deploy_cfg, common_cfg) runner.load_checkpoint(args.checkpoint) # image = Image.open(args.image) C, H, W = [int(_.strip()) for _ in args.dummy_input_shape.split(',')] dummy_image = np.random.random_integers(0, 255, (H, W, C)).astype(np.uint8) aug = runner.transform(image=dummy_image, label='') image, dummy_label = aug['image'], aug['label'] image = image.unsqueeze(0) input_len = runner.converter.test_encode(1)[0] model = runner.model need_text = runner.need_text if need_text: shape = tuple(image.shape), tuple(input_len.shape) else: shape = tuple(image.shape) dtypes = args.dtypes iters = args.iters int8_calibrator = None if args.calibration_images: calib_dataset = CalibDataset( args.calibration_images, runner.converter, runner.transform, need_text ) int8_calibrator = [ CALIBRATORS[mode](dataset=calib_dataset) for mode in args.calibration_modes ] if isinstance(runner.test_dataloader, dict): target_key = list(runner.test_dataloader.keys())[0] runner.logger.info( f'There are multiple datasets in for testing, using {target_key}' ) dataset = runner.test_dataloader[target_key].dataset else: dataset = runner.test_dataloader.dataset dataset = MetricDataset(dataset, runner.converter, need_text) metric = Metric(runner.metric, runner.converter) benchmark( model, shape, iters=iters, metric=metric, dtypes=dtypes, dataset=dataset, int8_calibrator=int8_calibrator, )