|
| 1 | +from typing import List, Dict, Any |
| 2 | +import torch |
| 3 | +import trtorch._C |
| 4 | +from trtorch._extra_info import _parse_extra_info |
| 5 | +from trtorch._version import __version__ |
| 6 | + |
| 7 | +def compile(module: torch.jit.ScriptModule, extra_info: Any) -> torch.jit.ScriptModule: |
| 8 | + """Compile a TorchScript module for NVIDIA GPUs using TensorRT |
| 9 | +
|
| 10 | + Takes a existing TorchScript module and a set of settings to configure the compiler |
| 11 | + and will convert methods to JIT Graphs which call equivalent TensorRT engines |
| 12 | +
|
| 13 | + Converts specifically the forward method of a TorchScript Module |
| 14 | +
|
| 15 | + Args: |
| 16 | + module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch |
| 17 | + ``torch.nn.Module`` |
| 18 | + extra_info (dict): Compilation settings including operating precision, target device, etc. |
| 19 | + One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs |
| 20 | + to the graph. All other keys are optional |
| 21 | +
|
| 22 | + .. code-block:: py |
| 23 | +
|
| 24 | + ExtraInfo = { |
| 25 | + "input_shapes": [ |
| 26 | + (1, 3, 224, 224), # Static input shape for input #1 |
| 27 | + { |
| 28 | + "min": (1, 3, 224, 224), |
| 29 | + "opt": (1, 3, 512, 512), |
| 30 | + "max": (1, 3, 1024, 1024) |
| 31 | + } # Dynamic input shape for input #2 |
| 32 | + ], |
| 33 | + "op_precision": torch.half, # Operating precision set to FP16 |
| 34 | + "refit": false, # enable refit |
| 35 | + "debug": false, # enable debuggable engine |
| 36 | + "strict_types": false, # kernels should strictly run in operating precision |
| 37 | + "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU |
| 38 | + "device": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA) |
| 39 | + "capability": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels |
| 40 | + "num_min_timing_iters": 2, # Number of minimization timing iterations used to select kernels |
| 41 | + "num_avg_timing_iters": 1, # Number of averaging timing iterations used to select kernels |
| 42 | + "workspace_size": 0, # Maximum size of workspace given to TensorRT |
| 43 | + "max_batch_size": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set) |
| 44 | + } |
| 45 | +
|
| 46 | + Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using |
| 47 | + torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum |
| 48 | + to select device type. |
| 49 | +
|
| 50 | + Returns: |
| 51 | + torch.jit.ScriptModule: Compiled TorchScript Module, when run it will execute via TensorRT |
| 52 | + """ |
| 53 | + compiled_cpp_mod = trtorch._C._compile_graph(module._c, _parse_extra_info(extra_info)) |
| 54 | + compiled_module = torch.jit._recursive.wrap_cpp_module(compiled_cpp_mod) |
| 55 | + return compiled_module |
| 56 | + |
| 57 | +def convert_method_to_trt_engine(module: torch.jit.ScriptModule, method_name: str, extra_info: Any) -> str: |
| 58 | + """Convert a TorchScript module method to a serialized TensorRT engine |
| 59 | +
|
| 60 | + Converts a specified method of a module to a serialized TensorRT engine given a dictionary of conversion settings |
| 61 | +
|
| 62 | + Args: |
| 63 | + module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch |
| 64 | + ``torch.nn.Module`` |
| 65 | + method_name (str): Name of method to convert |
| 66 | + extra_info (dict): Compilation settings including operating precision, target device, etc. |
| 67 | + One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs |
| 68 | + to the graph. All other keys are optional |
| 69 | +
|
| 70 | + .. code-block:: py |
| 71 | +
|
| 72 | + ExtraInfo = { |
| 73 | + "input_shapes": [ |
| 74 | + (1, 3, 224, 224), # Static input shape for input #1 |
| 75 | + { |
| 76 | + "min": (1, 3, 224, 224), |
| 77 | + "opt": (1, 3, 512, 512), |
| 78 | + "max": (1, 3, 1024, 1024) |
| 79 | + } # Dynamic input shape for input #2 |
| 80 | + ], |
| 81 | + "op_precision": torch.half, # Operating precision set to FP16 |
| 82 | + "refit": false, # enable refit |
| 83 | + "debug": false, # enable debuggable engine |
| 84 | + "strict_types": false, # kernels should strictly run in operating precision |
| 85 | + "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU |
| 86 | + "device": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA) |
| 87 | + "capability": trtorch.EngineCapability.DEFAULT, # Restrict kernel selection to safe gpu kernels or safe dla kernels |
| 88 | + "num_min_timing_iters": 2, # Number of minimization timing iterations used to select kernels |
| 89 | + "num_avg_timing_iters": 1, # Number of averaging timing iterations used to select kernels |
| 90 | + "workspace_size": 0, # Maximum size of workspace given to TensorRT |
| 91 | + "max_batch_size": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set) |
| 92 | + } |
| 93 | +
|
| 94 | + Input Sizes can be specified as torch sizes, tuples or lists. Op precisions can be specified using |
| 95 | + torch datatypes or trtorch datatypes and you can use either torch devices or the trtorch device type enum |
| 96 | + to select device type. |
| 97 | +
|
| 98 | + Returns: |
| 99 | + bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs |
| 100 | + """ |
| 101 | + return trtorch._C._convert_graph_to_trt_engine(module._c, method_name, _parse_extra_info(extra_info)) |
| 102 | + |
| 103 | +def check_method_op_support(module: torch.jit.ScriptModule, method_name: str) -> bool: |
| 104 | + """Checks to see if a method is fully supported by TRTorch |
| 105 | +
|
| 106 | + Checks if a method of a TorchScript module can be compiled by TRTorch, if not, a list of operators |
| 107 | + that are not supported are printed out and the function returns false, else true. |
| 108 | +
|
| 109 | + Args: |
| 110 | + module (torch.jit.ScriptModule): Source module, a result of tracing or scripting a PyTorch |
| 111 | + ``torch.nn.Module`` |
| 112 | + method_name (str): Name of method to check |
| 113 | +
|
| 114 | + Returns: |
| 115 | + bool: True if supported Method |
| 116 | + """ |
| 117 | + return trtorch._C._check_method_op_support(module._c, method_name) |
| 118 | + |
| 119 | +def dump_build_info(): |
| 120 | + """Prints build information about the TRTorch distribution to stdout |
| 121 | + """ |
| 122 | + print(get_build_info()) |
| 123 | + |
| 124 | +def get_build_info() -> str: |
| 125 | + """Returns a string containing the build information of TRTorch distribution |
| 126 | +
|
| 127 | + Returns: |
| 128 | + str: String containing the build information for TRTorch distribution |
| 129 | + """ |
| 130 | + build_info = trtorch._C._get_build_info() |
| 131 | + build_info = "TRTorch Version: " + str(__version__) + '\n' + build_info |
| 132 | + return build_info |
| 133 | + |
0 commit comments