예제 #1
0
def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)

    inference_cfg = cfg['inference']
    common_cfg = cfg.get('common')

    runner = InferenceRunner(inference_cfg, common_cfg)
    assert runner.use_gpu, 'Please use valid gpu to export model.'
    runner.load_checkpoint(args.checkpoint)
    model = runner.model

    shape = map(int, args.dummy_input_shape.split(','))
    dummy_input = torch.randn(1, *shape)

    if args.dynamic_shape:
        print(f'Convert to Onnx with dynamic input shape and '
              f'opset version {args.opset_version}')
    else:
        print(f'Convert to Onnx with constant input shape '
              f'{args.dummy_input_shape} and '
              f'opset version {args.opset_version}')
    torch2onnx(model,
               dummy_input,
               args.out,
               dynamic_shape=args.dynamic_shape,
               opset_version=args.opset_version,
               do_constant_folding=args.do_constant_folding,
               verbose=args.verbose)
    print(
        f'Convert successfully, saved onnx file: {os.path.abspath(args.out)}')
예제 #2
0
def main():
    args = parse_args()
    cfg_path = args.config
    cfg = Config.fromfile(cfg_path)

    multi_label = cfg.get('multi_label', False)
    inference_cfg = cfg['inference']
    common_cfg = cfg.get('common')

    runner = InferenceRunner(inference_cfg, common_cfg)
    runner.load_checkpoint(args.checkpoint)
    image = cv2.imread(args.image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, c = image.shape
    dummy_mask = np.zeros((h, w))
    output = runner(image, [dummy_mask])

    if args.need_resize:
        output = inverse_resize(output, image.shape)
    output = inverse_pad(output, image.shape)

    result(args.image,
           output,
           multi_label=multi_label,
           classes=CLASSES,
           palette=PALETTE,
           show=args.show,
           out=args.out)
예제 #3
0
def main():
    args = parse_args()
    out_name = args.out

    cfg_path = args.config
    cfg = Config.fromfile(cfg_path)

    inference_cfg = cfg['inference']
    common_cfg = cfg.get('common')

    runner = InferenceRunner(inference_cfg, common_cfg)
    assert runner.use_gpu, 'Please use valid gpu to export model.'
    runner.load_checkpoint(args.checkpoint)

    image = cv2.imread(args.image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, c = image.shape
    dummy_mask = np.zeros((h, w))
    image = runner.transform(image=image, masks=[dummy_mask])['image']

    dummy_input = image.unsqueeze(0).cuda()
    model = runner.model.cuda().eval()

    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.transform)
                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))
예제 #4
0
def main():
    args = parse_args()

    cfg_path = args.config
    cfg = Config.fromfile(cfg_path)

    test_cfg = cfg['test']
    inference_cfg = cfg['inference']
    base_cfg = cfg['common']

    runner = TestRunner(test_cfg, inference_cfg, base_cfg)
    assert runner.use_gpu, 'Please use gpu for benchmark.'
    runner.load_checkpoint(args.checkpoint)

    image = cv2.imread(args.image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    h, w, c = image.shape
    dummy_mask = np.zeros((h, w))
    image = runner.transform(image=image, masks=[dummy_mask])['image']

    dummy_input = image.unsqueeze(0).cuda()

    model = runner.model
    shape = tuple(dummy_input.shape)

    dtypes = args.dtypes
    iters = args.iters
    int8_calibrator = None
    if args.calibration_images:
        calib_dataset = CalibDataset(args.calibration_images, runner.transform)
        int8_calibrator = [
            CALIBRATORS[mode](dataset=calib_dataset)
            for mode in args.calibration_modes
        ]
    dataset = runner.test_dataloader.dataset
    metric = Metric(runner.metric, runner.compute)
    benchmark(model,
              shape,
              dtypes=dtypes,
              iters=iters,
              int8_calibrator=int8_calibrator,
              dataset=dataset,
              metric=metric)
예제 #5
0
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']
    inference_cfg = cfg['inference']
    common_cfg = cfg['common']
    common_cfg['workdir'] = workdir

    runner = TrainRunner(train_cfg, inference_cfg, common_cfg)
    runner()
예제 #6
0
파일: test.py 프로젝트: zfxu/vedaseg
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']
    inference_cfg = cfg['inference']
    common_cfg = cfg['common']
    common_cfg['workdir'] = workdir
    common_cfg['distribute'] = args.distribute

    runner = TestRunner(test_cfg, inference_cfg, common_cfg)
    runner.load_checkpoint(args.checkpoint)
    runner()