Skip to content

TensorRT

model_navigator.api.tensorrt

TensorRT optimize API.

optimize 

optimize(
    model,
    dataloader,
    sample_count=DEFAULT_SAMPLE_COUNT,
    batching=True,
    runners=None,
    optimization_profile=None,
    workspace=None,
    verbose=False,
    debug=False,
    verify_func=None,
)

Function executes correctness test, performance profling and optional verification on provided TensorRT model.

Parameters:

  • model (Union[Path, str]) –

    TensorRT model path or string

  • dataloader (SizedDataLoader) –

    Sized iterable with data that will be feed to the model

  • sample_count (Optional[int], default: DEFAULT_SAMPLE_COUNT ) –

    Limits how many samples will be used from dataloader

  • batching (Optional[bool], default: True ) –

    Enable or disable batching on first (index 0) dimension of the model

  • runners (Optional[Union[Union[str, Type[NavigatorRunner]], Tuple[Union[str, Type[NavigatorRunner]], ...]]], default: None ) –

    Use only runners provided as parameter

  • optimization_profile (Optional[OptimizationProfile], default: None ) –

    Optimization profile for conversion and profiling

  • workspace (Optional[Path], default: None ) –

    Workspace where packages will be extracted

  • verbose (bool, default: False ) –

    Enable verbose logging

  • debug (bool, default: False ) –

    Enable debug logging from commands

  • verify_func (Optional[VerifyFunction], default: None ) –

    Function for additional model verification

Returns:

  • Package –

    Package descriptor representing created package.

Source code in model_navigator/api/tensorrt.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def optimize(
    model: Union[Path, str],
    dataloader: SizedDataLoader,
    sample_count: Optional[int] = DEFAULT_SAMPLE_COUNT,
    batching: Optional[bool] = True,
    runners: Optional[Union[Union[str, Type[NavigatorRunner]], Tuple[Union[str, Type[NavigatorRunner]], ...]]] = None,
    optimization_profile: Optional[OptimizationProfile] = None,
    workspace: Optional[Path] = None,
    verbose: bool = False,
    debug: bool = False,
    verify_func: Optional[VerifyFunction] = None,
) -> Package:
    """Function executes correctness test, performance profling and optional verification on provided TensorRT model.

    Args:
        model: TensorRT model path or string
        dataloader: Sized iterable with data that will be feed to the model
        sample_count: Limits how many samples will be used from dataloader
        batching: Enable or disable batching on first (index 0) dimension of the model
        runners: Use only runners provided as parameter
        optimization_profile: Optimization profile for conversion and profiling
        workspace: Workspace where packages will be extracted
        verbose: Enable verbose logging
        debug: Enable debug logging from commands
        verify_func: Function for additional model verification

    Returns:
        Package descriptor representing created package.
    """
    if isinstance(model, str):
        model = Path(model)
    target_formats = DEFAULT_TENSORRT_TARGET_FORMATS
    target_device = DeviceKind.CUDA

    if runners is None:
        runners = default_runners(device_kind=target_device)

    if optimization_profile is None:
        optimization_profile = OptimizationProfile()

    runner_names = enums.parse(runners, lambda runner: runner if isinstance(runner, str) else runner.name())

    config = CommonConfig(
        Framework.TENSORRT,
        model=model,
        dataloader=dataloader,
        target_formats=target_formats,
        target_device=target_device,
        sample_count=sample_count,
        batch_dim=0 if batching else None,
        runner_names=runner_names,
        optimization_profile=optimization_profile,
        verbose=verbose,
        debug=debug,
        verify_func=verify_func,
        custom_configs=None,
    )

    models_config = ModelConfigBuilder.generate_model_config(
        framework=Framework.TENSORRT,
        target_formats=target_formats,
        custom_configs=None,
    )

    builders = [
        preprocessing_builder,
        find_device_max_batch_size_builder,
        correctness_builder,
        performance_builder,
        verify_builder,
    ]

    package = optimize_pipeline(
        model=model,
        workspace=workspace,
        builders=builders,
        config=config,
        models_config=models_config,
    )

    return package