| | from dataclasses import dataclass |
| | from typing import Dict, List, Optional |
| |
|
| | import numpy as np |
| | import tensorrt as trt |
| | import torch |
| |
|
| | from tensorrt_llm._utils import trt_dtype_to_str, trt_dtype_to_torch |
| | from tensorrt_llm.logger import logger |
| | from tensorrt_llm.network import Network, get_plugin_info, set_plugin_info |
| | from tensorrt_llm.plugin.plugin import PluginConfig |
| | from tensorrt_llm.runtime.session import Session |
| |
|
| | from .utils import (current_flags, get_builder_flags, get_sorted_layer_ids, |
| | get_strongly_typed, get_trt_network, set_trt_network, |
| | to_base_class_layer, to_subclass_layer) |
| |
|
| |
|
| | class Tensor: |
| |
|
| | def __init__(self, graph: "PipelineGraph"): |
| | self._graph = graph |
| | self._trt = None |
| | self._shape = None |
| | self._max_shape = None |
| | self._value = None |
| | self.producer: Layer = None |
| | self.output_index = None |
| | self.consumers = [] |
| | self.graph_input_index = -1 |
| | self.graph_output_index = -1 |
| | self.attrs = {} |
| |
|
| | @staticmethod |
| | def from_trt(graph: "PipelineGraph", trt_tensor: trt.ITensor): |
| | tensor = Tensor(graph) |
| | tensor._trt = trt_tensor |
| | return tensor |
| |
|
| | def as_trt(self) -> trt.ITensor: |
| | return self._trt |
| |
|
| | def copy(self) -> "Tensor": |
| | tensor = Tensor(self._graph) |
| | tensor._trt = self._trt |
| | tensor._shape = self._shape |
| | tensor._max_shape = self._max_shape |
| | tensor._value = self._value |
| | tensor.producer = self.producer |
| | tensor.output_index = self.output_index |
| | tensor.consumers = [*self.consumers] |
| | tensor.graph_input_index = self.graph_input_index |
| | tensor.graph_output_index = self.graph_output_index |
| | tensor.attrs = self.attrs.copy() |
| | return tensor |
| |
|
| | @property |
| | def graph(self) -> "PipelineGraph": |
| | return self._graph |
| |
|
| | @property |
| | def name(self) -> str: |
| | return self._trt.name |
| |
|
| | @name.setter |
| | def name(self, name: str): |
| | old_name = self._trt.name |
| | if name != old_name: |
| | self._trt.name = name |
| | self.graph._tensors[name] = self |
| | del self.graph._tensors[old_name] |
| | if self.is_graph_input: |
| | self.graph._inputs[name] = self |
| | del self.graph._inputs[old_name] |
| | elif self.is_graph_output: |
| | self.graph._outputs[name] = self |
| | del self.graph._outputs[old_name] |
| |
|
| | @property |
| | def shape(self): |
| | return self._shape |
| |
|
| | @property |
| | def max_shape(self): |
| | return self._max_shape |
| |
|
| | @property |
| | def raw_shape(self): |
| | assert isinstance(self._trt, trt.ITensor) |
| | return self._trt.shape |
| |
|
| | @shape.setter |
| | def shape(self, shape): |
| | self._shape = shape |
| |
|
| | @max_shape.setter |
| | def max_shape(self, max_shape): |
| | self._max_shape = max_shape |
| |
|
| | @raw_shape.setter |
| | def raw_shape(self, raw_shape): |
| | assert isinstance(self._trt, trt.ITensor) |
| | self._trt.shape = raw_shape |
| |
|
| | @property |
| | def value(self): |
| | return self._value |
| |
|
| | @value.setter |
| | def value(self, value): |
| | self._value = value |
| |
|
| | @property |
| | def dtype(self): |
| | return self._trt.dtype |
| |
|
| | @property |
| | def broadcast_across_batch(self): |
| | return self._trt.broadcast_across_batch |
| |
|
| | @property |
| | def dtype_size(self): |
| | return self.dtype.itemsize |
| |
|
| | @property |
| | def dtype_str(self): |
| | return trt_dtype_to_str(self.dtype) |
| |
|
| | @property |
| | def dtype_str_size(self): |
| | return [trt_dtype_to_str(self.dtype), self.dtype.itemsize] |
| |
|
| | @property |
| | def is_graph_input(self) -> bool: |
| | return self.graph_input_index != -1 |
| |
|
| | @property |
| | def is_graph_output(self) -> bool: |
| | return self.graph_output_index != -1 |
| |
|
| | @property |
| | def is_graph_io(self) -> bool: |
| | return self.is_graph_input or self.is_graph_output |
| |
|
| |
|
| | class Layer: |
| |
|
| | def __init__(self, graph): |
| | self._graph = graph |
| | self._trt = None |
| | self._index = None |
| | self._inputs = [] |
| | self._outputs = [] |
| | self._is_shape_io = False |
| | self.attrs = {} |
| |
|
| | @staticmethod |
| | def from_trt(graph, trt_layer, index): |
| | layer = Layer(graph) |
| | layer._trt = trt_layer |
| | layer._index = index |
| | for i in range(trt_layer.num_inputs): |
| | input = trt_layer.get_input(i) |
| | if input is not None: |
| | layer._inputs.append(graph.get_tensor(input.name)) |
| | layer._inputs[i].consumers.append((layer, i)) |
| | else: |
| | layer._inputs.append(None) |
| | for i in range(trt_layer.num_outputs): |
| | output = trt_layer.get_output(i) |
| | layer._outputs.append(graph.get_tensor(output.name)) |
| | layer._outputs[i].producer = layer |
| | layer._outputs[i].output_index = i |
| | set_trt_network(trt_layer, graph.as_trt()) |
| | return layer |
| |
|
| | def as_trt(self) -> trt.ILayer: |
| | return self._trt |
| |
|
| | @property |
| | def graph(self) -> "PipelineGraph": |
| | return self._graph |
| |
|
| | @property |
| | def name(self) -> str: |
| | return self._trt.name |
| |
|
| | @name.setter |
| | def name(self, name: str): |
| | old_name = self._trt.name |
| | if name != old_name: |
| | self._trt.name = name |
| | self.graph._layers[name] = self |
| | del self.graph._layers[old_name] |
| |
|
| | @property |
| | def type(self) -> trt.LayerType: |
| | return self._trt.type |
| |
|
| | @property |
| | def index(self) -> int: |
| | return self._index |
| |
|
| | @property |
| | def inputs(self) -> List[Tensor]: |
| | return self._inputs |
| |
|
| | @property |
| | def outputs(self) -> List[Tensor]: |
| | return self._outputs |
| |
|
| | def get_input(self, index: int) -> Tensor: |
| | return self._inputs[index] |
| |
|
| | def get_output(self, index: int) -> Tensor: |
| | return self._outputs[index] |
| |
|
| | @property |
| | def num_inputs(self) -> int: |
| | return self._trt.num_inputs |
| |
|
| | @property |
| | def num_outputs(self) -> int: |
| | return self._trt.num_outputs |
| |
|
| | @property |
| | def is_shape_io(self) -> bool: |
| | return self._is_shape_io |
| |
|
| | def to_subclass(self): |
| | to_subclass_layer(self._trt) |
| |
|
| | def to_base_class(self): |
| | to_base_class_layer(self._trt) |
| |
|
| | def assign_shapes(self, shapes, values): |
| | for output in self.outputs: |
| | output.shape = shapes[output.name] |
| | output.value = values.get(output.name) |
| |
|
| |
|
| | @dataclass |
| | class GraphRunner: |
| | session: Session |
| | inputs: Dict[str, torch.Tensor] |
| | outputs: Dict[str, torch.Tensor] |
| | stream: torch.Stream |
| |
|
| | def run(self): |
| | cuda_stream = self.stream.cuda_stream |
| | assert self.session.run(self.inputs, self.outputs, cuda_stream) |
| | self.stream.synchronize() |
| | return self.outputs |
| |
|
| |
|
| | class PipelineGraph: |
| |
|
| | def __init__(self): |
| | self._trt = None |
| | self._inputs: Dict[str, Tensor] = {} |
| | self._outputs: Dict[str, Tensor] = {} |
| | self._layers: Dict[str, Layer] = {} |
| | self._tensors: Dict[str, Tensor] = {} |
| | self._io_buffer_mapping = {} |
| | self._unfilled_weights = {} |
| | self._auto_parallel_config = None |
| | self._plugin_config: PluginConfig = None |
| |
|
| | @staticmethod |
| | def create_graph(): |
| | graph = PipelineGraph() |
| | trt_builder = trt.Builder(logger.trt_logger) |
| | explicit_batch_flag = 0 |
| | |
| | if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys( |
| | ): |
| | explicit_batch_flag = 1 << int( |
| | trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) |
| | if get_strongly_typed(): |
| | network = trt_builder.create_network( |
| | explicit_batch_flag |
| | | (1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))) |
| | else: |
| | network = trt_builder.create_network(explicit_batch_flag) |
| | graph._trt = network |
| | return graph |
| |
|
| | def _register_unfilled_weights(self, layer_name, weights, values): |
| | self._unfilled_weights[layer_name] = (weights, values) |
| |
|
| | def _add_tensor(self, tensor, old_tensor, prefix): |
| | if prefix is not None: |
| | tensor.name = prefix + old_tensor.name |
| | else: |
| | tensor.name = old_tensor.name |
| | tensor.location = old_tensor.location |
| | if old_tensor.dynamic_range is not None: |
| | tensor.dynamic_range = old_tensor.dynamic_range |
| | if tensor.is_network_input: |
| | tensor.shape = old_tensor.shape |
| | for i in range(len(old_tensor.shape)): |
| | name = old_tensor.get_dimension_name(i) |
| | if name is not None: |
| | tensor.set_dimension_name(i, name) |
| | return self._register_tensor(tensor) |
| |
|
| | def _register_tensor(self, tensor): |
| | wrapped_tensor = Tensor.from_trt(self, tensor) |
| | assert tensor.name not in self._tensors |
| | self._tensors[tensor.name] = wrapped_tensor |
| | return wrapped_tensor |
| |
|
| | def add_input(self, tensor, prefix=None): |
| | tensor_name = tensor.name |
| | if prefix is not None: |
| | tensor_name = prefix + tensor_name |
| | input = self._trt.add_input(tensor_name, tensor.dtype, tensor.shape) |
| | new_tensor = self._add_tensor(input, tensor, prefix) |
| | new_tensor.graph_input_index = len(self._inputs) |
| | self._inputs[tensor_name] = new_tensor |
| | return new_tensor |
| |
|
| | def register_input(self, tensor, index=None): |
| | if index is None: |
| | index = self.num_inputs - 1 |
| | assert self._trt.get_input(index).name == tensor.name |
| | wrapped_input = self._register_tensor(tensor) |
| | wrapped_input.graph_input_index = index |
| | self._inputs[tensor.name] = wrapped_input |
| | return wrapped_input |
| |
|
| | def add_output(self, tensor, prefix=None): |
| | tensor_name = tensor.name |
| | if prefix is not None: |
| | tensor_name = prefix + tensor_name |
| | output = self.get_tensor(tensor_name) |
| | output.graph_output_index = len(self._outputs) |
| | trt_output = output.as_trt() |
| | self._trt.mark_output(trt_output) |
| | trt_output.dtype = tensor.dtype |
| | self._outputs[tensor_name] = output |
| | return output |
| |
|
| | def add_output_shape(self, tensor, prefix=None): |
| | tensor_name = tensor.name |
| | if prefix is not None: |
| | tensor_name = prefix + tensor_name |
| | output = self.get_tensor(tensor_name) |
| | trt_output = output.as_trt() |
| | self._trt.mark_output_for_shapes(trt_output) |
| | trt_output.dtype = tensor.dtype |
| | self._outputs[tensor_name] = output |
| | return output |
| |
|
| | def add_layer( |
| | self, |
| | layer, |
| | input_mapping=None, |
| | prefix=None, |
| | updated_attrs=None, |
| | ) -> Layer: |
| |
|
| | def get_input(i): |
| | name = layer.get_input(i).name |
| | if prefix is not None: |
| | name = prefix + name |
| | if input_mapping is not None and name in input_mapping: |
| | name = input_mapping[name] |
| | return self.get_tensor(name).as_trt() |
| |
|
| | network = self._trt |
| | layer_type = layer.type |
| | to_subclass_layer(layer) |
| | if layer_type == trt.LayerType.ACTIVATION: |
| | trt_input = get_input(0) |
| | new_layer = network.add_activation(trt_input, layer.type) |
| | new_layer.alpha = layer.alpha |
| | new_layer.beta = layer.beta |
| | elif layer_type == trt.LayerType.CONCATENATION: |
| | trt_inputs = [get_input(i) for i in range(layer.num_inputs)] |
| | new_layer = network.add_concatenation(trt_inputs) |
| | new_layer.axis = layer.axis |
| | elif layer_type == trt.LayerType.CONSTANT: |
| | new_layer = network.add_constant(layer.shape, layer.weights) |
| | elif layer_type == trt.LayerType.ELEMENTWISE: |
| | new_layer = network.add_elementwise(get_input(0), get_input(1), |
| | layer.op) |
| | elif layer_type == trt.LayerType.FILL: |
| | if layer.num_inputs >= 1 and layer.get_input(0) is not None: |
| | shape_input = get_input(0) |
| | shape = [1] |
| | else: |
| | shape_input = None |
| | shape = layer.shape |
| | new_layer = network.add_fill(shape, layer.operation, layer.to_type) |
| | if shape_input is not None: |
| | new_layer.set_input(0, shape_input) |
| | if layer.num_inputs >= 1 and layer.get_input(0) is not None: |
| | new_layer.set_input(0, get_input(0)) |
| | if layer.num_inputs >= 2 and layer.get_input(1) is not None: |
| | new_layer.set_input(1, get_input(1)) |
| | else: |
| | new_layer.alpha = layer.alpha |
| | if layer.num_inputs >= 3 and layer.get_input(2) is not None: |
| | new_layer.set_input(2, get_input(2)) |
| | else: |
| | new_layer.beta = layer.beta |
| | elif layer_type == trt.LayerType.GATHER: |
| | trt_input = get_input(0) |
| | trt_indices = get_input(1) |
| | new_layer = network.add_gather_v2(trt_input, trt_indices, |
| | layer.mode) |
| | new_layer.axis = layer.axis |
| | new_layer.num_elementwise_dims = layer.num_elementwise_dims |
| | new_layer.mode = layer.mode |
| | elif layer_type == trt.LayerType.MATRIX_MULTIPLY: |
| | new_layer = network.add_matrix_multiply(get_input(0), layer.op0, |
| | get_input(1), layer.op1) |
| | elif layer_type == trt.LayerType.REDUCE: |
| | new_layer = network.add_reduce(get_input(0), layer.op, layer.axes, |
| | layer.keep_dims) |
| | elif layer_type == trt.LayerType.SELECT: |
| | trt_condition = get_input(0) |
| | trt_then = get_input(1) |
| | trt_else = get_input(2) |
| | new_layer = network.add_select(trt_condition, trt_then, trt_else) |
| | elif layer_type == trt.LayerType.SHUFFLE: |
| | new_layer = network.add_shuffle(get_input(0)) |
| | new_layer.first_transpose = layer.first_transpose |
| | new_layer.second_transpose = layer.second_transpose |
| | new_layer.zero_is_placeholder = layer.zero_is_placeholder |
| | if layer.num_inputs >= 2: |
| | trt_reshape_dims_tensor = get_input(1) |
| | new_layer.set_input(1, trt_reshape_dims_tensor) |
| | else: |
| | new_layer.reshape_dims = layer.reshape_dims |
| | elif layer_type == trt.LayerType.SLICE: |
| | if layer.num_inputs >= 2 and layer.get_input(1) is not None: |
| | trt_start = get_input(1) |
| | start = [] |
| | else: |
| | trt_start = None |
| | start = layer.start |
| | if layer.num_inputs >= 3 and layer.get_input(2) is not None: |
| | trt_shape = get_input(2) |
| | shape = [] |
| | else: |
| | trt_shape = None |
| | shape = layer.shape |
| | if layer.num_inputs >= 4 and layer.get_input(3) is not None: |
| | trt_stride = get_input(3) |
| | stride = [] |
| | else: |
| | trt_stride = None |
| | stride = layer.stride |
| | new_layer = network.add_slice(get_input(0), start, shape, stride) |
| | new_layer.mode = layer.mode |
| | if trt_start is not None: |
| | new_layer.set_input(1, trt_start) |
| | if trt_shape is not None: |
| | new_layer.set_input(2, trt_shape) |
| | if trt_stride is not None: |
| | new_layer.set_input(3, trt_stride) |
| | elif layer_type == trt.LayerType.SOFTMAX: |
| | new_layer = network.add_softmax(get_input(0)) |
| | new_layer.axes = layer.axes |
| | elif layer_type == trt.LayerType.UNARY: |
| | new_layer = network.add_unary(get_input(0), layer.op) |
| | elif layer_type == trt.LayerType.SHAPE: |
| | new_layer = network.add_shape(get_input(0)) |
| | elif layer_type == trt.LayerType.ASSERTION: |
| | new_layer = network.add_assertion(get_input(0), layer.message) |
| | elif layer_type == trt.LayerType.CAST: |
| | new_layer = network.add_cast(get_input(0), layer.to_type) |
| | elif layer_type == trt.LayerType.NORMALIZATION: |
| | trt_input = get_input(0) |
| | trt_scale = get_input(1) |
| | trt_bias = get_input(2) |
| | new_layer = network.add_normalization(trt_input, trt_scale, |
| | trt_bias, layer.axes) |
| | new_layer.epsilon = layer.epsilon |
| | new_layer.num_groups = layer.num_groups |
| | new_layer.compute_precision = layer.compute_precision |
| | elif layer_type == trt.LayerType.IDENTITY: |
| | new_layer = network.add_identity(get_input(0)) |
| | elif layer_type == trt.LayerType.PLUGIN_V2: |
| | plugin = layer.plugin |
| | updated = False |
| | if (updated_attrs is not None |
| | and updated_attrs.get("plugin") is not None): |
| | plugin = updated_attrs["plugin"] |
| | updated = True |
| | updated_attrs = None |
| | new_layer = network.add_plugin_v2( |
| | [get_input(i) for i in range(layer.num_inputs)], |
| | plugin, |
| | ) |
| | else: |
| | raise NotImplementedError( |
| | "Unsupported layer type: {}".format(layer_type)) |
| |
|
| | if updated_attrs is not None: |
| | for attr_name, attr_value in updated_attrs.items(): |
| | setattr(new_layer, attr_name, attr_value) |
| |
|
| | to_base_class_layer(layer) |
| | to_base_class_layer(new_layer) |
| | layer_index = network.num_layers - 1 |
| | layer_name = layer.name |
| | if prefix is not None: |
| | layer_name = prefix + layer_name |
| | new_layer.name = layer_name |
| | new_layer.metadata = new_layer.name |
| | if layer.precision_is_set: |
| | new_layer.precision = layer.precision |
| | for i in range(layer.num_outputs): |
| | if layer.output_type_is_set(i): |
| | new_layer.set_output_type(i, layer.get_output_type(i)) |
| | output = new_layer.get_output(i) |
| | self._add_tensor(output, layer.get_output(i), prefix) |
| | wrapped_layer = Layer.from_trt(self, new_layer, layer_index) |
| | assert layer_name not in self._layers |
| | self._layers[layer_name] = wrapped_layer |
| | if layer_type == trt.LayerType.PLUGIN_V2: |
| | if not updated: |
| | plugin_info = get_plugin_info(get_trt_network(layer), |
| | layer.name) |
| | set_plugin_info(self.as_trt(), new_layer.name, plugin_info) |
| | return wrapped_layer |
| |
|
| | def register_layer(self, layer, index=None): |
| | if index is None: |
| | index = self.num_layers - 1 |
| | assert self._trt.get_layer(index).name == layer.name |
| | to_base_class_layer(layer) |
| | for i in range(layer.num_outputs): |
| | output = layer.get_output(i) |
| | self._register_tensor(output) |
| | wrapped_layer = Layer.from_trt(self, layer, index) |
| | assert layer.name not in self._layers |
| | self._layers[layer.name] = wrapped_layer |
| | to_subclass_layer(layer) |
| | return wrapped_layer |
| |
|
| | def get_runner( |
| | self, |
| | shapes=None, |
| | values=None, |
| | profile=None, |
| | timing_cache=None, |
| | opt_level=None, |
| | ) -> GraphRunner: |
| | shapes = shapes or {} |
| | values = values or {} |
| | inputs = {} |
| | outputs = {} |
| | for input in self.inputs: |
| | if input is not None: |
| | value = values.get(input.name) |
| | if value is None: |
| | value = input.value |
| | if value is not None: |
| | if not isinstance(value, torch.Tensor): |
| | value = torch.tensor( |
| | value, |
| | dtype=trt_dtype_to_torch(input.dtype), |
| | device='cpu', |
| | ) |
| | inputs[input.name] = value |
| | else: |
| | shape = shapes.get(input.name) |
| | if shape is None: |
| | shape = input.shape |
| | assert shape is not None |
| | inputs[input.name] = torch.empty( |
| | tuple(shape), |
| | dtype=trt_dtype_to_torch(input.dtype), |
| | device=torch.cuda.current_device(), |
| | ) |
| | if torch.is_floating_point(inputs[input.name]): |
| | inputs[input.name].normal_() |
| | |
| | for output in self.outputs: |
| | if output.as_trt().is_shape_tensor: |
| | continue |
| | if output.name in self._io_buffer_mapping: |
| | input_name = self._io_buffer_mapping[output.name] |
| | if input_name in inputs: |
| | outputs[output.name] = inputs[input_name] |
| | continue |
| | value = values.get(output.name) |
| | if value is not None and isinstance(value, torch.Tensor): |
| | outputs[output.name] = value |
| | else: |
| | shape = shapes.get(output.name) |
| | if shape is None: |
| | shape = output.shape |
| | assert shape is not None |
| | outputs[output.name] = torch.empty( |
| | tuple(shape), |
| | dtype=trt_dtype_to_torch(output.dtype), |
| | device=torch.cuda.current_device(), |
| | ) |
| | network = self.as_trt() |
| | config = network.builder.create_builder_config() |
| | if opt_level is not None: |
| | config.builder_optimization_level = opt_level |
| | config.flags = get_builder_flags() |
| | profile = profile or network.builder.create_optimization_profile() |
| | profile_index = config.add_optimization_profile(profile) |
| | if timing_cache is not None: |
| | config.set_timing_cache(timing_cache, ignore_mismatch=False) |
| | plan = network.builder.build_serialized_network(network, config) |
| | if plan is None: |
| | logger.error('Engine building failed, please check the error log.') |
| | session = Session.from_serialized_engine(plan) |
| | stream = torch.cuda.current_stream() |
| | cuda_stream = stream.cuda_stream |
| | context = session.context |
| | context.set_optimization_profile_async(profile_index, cuda_stream) |
| | runner = GraphRunner(session, inputs, outputs, stream) |
| | return runner |
| |
|
| | def run( |
| | self, |
| | shapes=None, |
| | values=None, |
| | profile=None, |
| | timing_cache=None, |
| | opt_level=None, |
| | ): |
| | return self.get_runner( |
| | shapes, |
| | values, |
| | profile, |
| | timing_cache, |
| | opt_level, |
| | ).run() |
| |
|
| | def duplicate_graph(self): |
| | graph = PipelineGraph.create_graph() |
| | network = self.as_trt() |
| | for i in range(network.num_inputs): |
| | input = network.get_input(i) |
| | graph.add_input(input) |
| | sorted_layer_ids = get_sorted_layer_ids(network) |
| | for i in sorted_layer_ids: |
| | layer = network.get_layer(i) |
| | graph.add_layer(layer) |
| | for i in range(network.num_outputs): |
| | output = network.get_output(i) |
| | if output.is_shape_tensor: |
| | graph.add_output_shape(output) |
| | else: |
| | graph.add_output(output) |
| | return graph |
| |
|
| | @staticmethod |
| | def from_trt(trt_network): |
| | graph = PipelineGraph() |
| | graph._trt = trt_network |
| |
|
| | |
| | for i in range(trt_network.num_inputs): |
| | trt_input = trt_network.get_input(i) |
| | tensor = Tensor.from_trt(graph, trt_input) |
| | tensor.graph_input_index = i |
| | graph._tensors[tensor.name] = tensor |
| | graph._inputs[tensor.name] = tensor |
| | for i in range(trt_network.num_layers): |
| | trt_layer = trt_network.get_layer(i) |
| | for i in range(trt_layer.num_outputs): |
| | trt_output = trt_layer.get_output(i) |
| | tensor = Tensor.from_trt(graph, trt_output) |
| | graph._tensors[tensor.name] = tensor |
| |
|
| | |
| | for i in range(trt_network.num_layers): |
| | layer = Layer.from_trt(graph, trt_network.get_layer(i), i) |
| | graph._layers[layer.name] = layer |
| | for i in range(trt_network.num_outputs): |
| | tensor_name = trt_network.get_output(i).name |
| | output_tensor = graph._tensors[tensor_name] |
| | output_tensor.graph_output_index = i |
| | graph._outputs[tensor_name] = output_tensor |
| |
|
| | return graph |
| |
|
| | @staticmethod |
| | def from_network(network: Network, builder_config): |
| | builder_flags = builder_config.trt_builder_config.flags |
| | with current_flags(builder_flags, network.strongly_typed): |
| | graph = PipelineGraph.from_trt(network.trt_network) |
| | graph.infer_shapes(network._generate_optimization_profiles()[-1]) |
| | return graph |
| |
|
| | def assign_shapes(self, shape_info=None, is_partial=False): |
| | if shape_info is None: |
| | for tensor in self.tensors: |
| | tensor.shape = tensor.raw_shape |
| | return |
| | for tensor in self.tensors: |
| | if tensor.name in shape_info.shapes: |
| | tensor.shape = shape_info.shapes[tensor.name] |
| | elif not is_partial: |
| | raise ValueError(f"Cannot find shape for tensor: {tensor.name}") |
| | if shape_info.max_shapes is not None: |
| | if tensor.name in shape_info.max_shapes: |
| | tensor.max_shape = shape_info.max_shapes[tensor.name] |
| | elif not is_partial: |
| | raise ValueError( |
| | f"Cannot find max shape for tensor: {tensor.name}") |
| | if tensor.name in shape_info.values: |
| | tensor.value = shape_info.values[tensor.name] |
| | for layer in self.layers: |
| | if layer.name in shape_info.shape_layers: |
| | layer._is_shape_io = True |
| |
|
| | def infer_shapes(self, profile=None): |
| | from .shape_info import get_shape_info |
| |
|
| | shape_info = get_shape_info(self._trt, profile) |
| | self.assign_shapes(shape_info) |
| |
|
| | def as_trt(self) -> trt.INetworkDefinition: |
| | return self._trt |
| |
|
| | def get_input(self, name: str) -> Tensor: |
| | return self._inputs.get(name) |
| |
|
| | def is_input(self, name: str) -> bool: |
| | return name in self._inputs |
| |
|
| | @property |
| | def inputs(self) -> List[Tensor]: |
| | return [*self._inputs.values()] |
| |
|
| | @property |
| | def num_inputs(self) -> int: |
| | return self._trt.num_inputs |
| |
|
| | def get_output(self, name: str) -> Tensor: |
| | return self._outputs.get(name) |
| |
|
| | def is_output(self, name: str) -> bool: |
| | return name in self._outputs |
| |
|
| | @property |
| | def outputs(self) -> List[Tensor]: |
| | return [*self._outputs.values()] |
| |
|
| | @property |
| | def num_outputs(self) -> int: |
| | return self._trt.num_outputs |
| |
|
| | def get_tensor(self, name: str) -> Tensor: |
| | return self._tensors.get(name) |
| |
|
| | @property |
| | def tensors(self) -> List[Tensor]: |
| | return [*self._tensors.values()] |
| |
|
| | def get_layer(self, name: str) -> Layer: |
| | return self._layers.get(name) |
| |
|
| | @property |
| | def layers(self) -> List[Layer]: |
| | return [*self._layers.values()] |
| |
|
| | @property |
| | def sorted_layers(self) -> List[Layer]: |
| | sorted_layer_ids = get_sorted_layer_ids(self.as_trt()) |
| | return [ |
| | self.get_layer(self.as_trt().get_layer(layer_id).name) |
| | for layer_id in sorted_layer_ids |
| | ] |
| |
|
| | @property |
| | def num_layers(self) -> int: |
| | return self._trt.num_layers |
| |
|
| | def to_dot(self, |
| | path=None, |
| | per_device=False, |
| | per_block=False, |
| | ignore_shape_io=False, |
| | no_style=False, |
| | extra_attrs=None) -> Optional[str]: |
| | ''' |
| | Get a graphviz representation of the graph. |
| | |
| | Parameters: |
| | path: the path to save the graphviz file, if not provided, will return the graphviz source code |
| | ''' |
| | try: |
| | import graphviz |
| | except ImportError: |
| | logger.error( |
| | "Failed to import graphviz, please install graphviz to enable PipelineGraph.to_dot()" |
| | ) |
| | return |
| |
|
| | extra_attrs = extra_attrs or [] |
| |
|
| | graph = graphviz.Digraph() |
| | input_block_graph = graphviz.Digraph(name='cluster_inputs') |
| | output_block_graph = graphviz.Digraph(name='cluster_outputs') |
| | device_graphs = {} |
| | block_graphs = {} |
| | block_graph_mapping = [] |
| | tensor_names = set() |
| | layer_names = set() |
| |
|
| | common_style = dict(fontname='Arial', ) |
| | node_style = dict( |
| | **common_style, |
| | style='rounded,filled,bold', |
| | ) |
| | tensor_style = dict( |
| | **node_style, |
| | shape='ellipse', |
| | fillcolor='white', |
| | ) |
| | input_tensor_style = {**tensor_style, 'fillcolor': 'green'} |
| | output_tensor_style = {**tensor_style, 'fillcolor': 'lightgreen'} |
| | layer_style = dict( |
| | **node_style, |
| | shape='box', |
| | fillcolor='white', |
| | ) |
| | shape_layer_style = {**layer_style, 'fillcolor': 'grey'} |
| | helper_layer_style = {**layer_style, 'fillcolor': 'lightgrey'} |
| | graph_style = dict( |
| | **common_style, |
| | style='rounded', |
| | penwidth='5', |
| | fontsize='28', |
| | ) |
| | device_graph_style = dict( |
| | **graph_style, |
| | color='cornflowerblue', |
| | ) |
| | block_graph_style = dict( |
| | **graph_style, |
| | color='darkcyan', |
| | ) |
| | input_block_style = dict( |
| | **graph_style, |
| | color='green', |
| | ) |
| | output_block_style = dict( |
| | **graph_style, |
| | color='lightgreen', |
| | ) |
| | if no_style: |
| | device_graph_style = {} |
| | block_graph_style = {} |
| | input_block_style = {} |
| | output_block_style = {} |
| | input_block_graph.attr(label='inputs', **input_block_style) |
| | output_block_graph.attr(label='outputs', **output_block_style) |
| |
|
| | def get_tensor_labels(tensor): |
| | labels = [] |
| | if tensor.value is not None: |
| | labels.append(f"value={tensor.value}") |
| | else: |
| | labels.append(f"dtype={tensor.dtype.name}{tensor.shape}") |
| | for attr_name in extra_attrs: |
| | if attr_name in tensor.attrs: |
| | labels.append(f"{attr_name}={tensor.attrs[attr_name]}") |
| | return labels |
| |
|
| | def get_device_graph(name): |
| | if per_device and name.startswith('device'): |
| | device_name = name.split('_')[0] |
| | if device_name not in device_graphs: |
| | device_graph = graphviz.Digraph(name='cluster_' + |
| | device_name) |
| | device_graph.attr(label=device_name, **device_graph_style) |
| | device_graphs[device_name] = device_graph |
| | return device_graphs[device_name] |
| | return None |
| |
|
| | def get_block_graph(layer, current_graph): |
| | if per_block and 'block_id' in layer.attrs: |
| | block_label = f"block{layer.attrs['block_id']}" |
| | if current_graph.name is not None: |
| | graph_label = current_graph.name[len('cluster_'):] |
| | else: |
| | graph_label = '' |
| | block_name = f"{graph_label}{block_label}" |
| | if block_name not in block_graphs: |
| | block_graph = graphviz.Digraph(name='cluster_' + block_name) |
| | block_graph.attr(label=block_label, **block_graph_style) |
| | block_graphs[block_name] = block_graph |
| | block_graph_mapping.append((current_graph, block_graph)) |
| | return block_graphs[block_name] |
| | return current_graph |
| |
|
| | for name, tensor in self._tensors.items(): |
| | style = tensor_style |
| | if tensor.is_graph_input: |
| | style = input_tensor_style |
| | current_graph = input_block_graph |
| | elif tensor.is_graph_output: |
| | style = output_tensor_style |
| | current_graph = output_block_graph |
| | elif tensor.producer.num_outputs == 1: |
| | continue |
| | else: |
| | current_graph = get_device_graph(name) or graph |
| | current_graph = get_block_graph(tensor.producer, current_graph) |
| | if no_style: |
| | style = {} |
| | labels = [name, *get_tensor_labels(tensor)] |
| | content = "\n".join(labels) |
| | current_graph.node(name, content, **style) |
| | tensor_names.add(name) |
| |
|
| | for layer in self.sorted_layers: |
| | name = layer.name |
| |
|
| | style = layer_style |
| | if layer.is_shape_io: |
| | if ignore_shape_io: |
| | continue |
| | style = shape_layer_style |
| | elif layer.attrs.get("role", None) == "helper": |
| | style = helper_layer_style |
| | fillcolor = None |
| | plugin_type = None |
| | if layer.type == trt.LayerType.PLUGIN_V2: |
| | fillcolor = 'yellow' |
| | layer.to_subclass() |
| | plugin_type = layer.as_trt().plugin.plugin_type |
| | layer.to_base_class() |
| | if layer.type == trt.LayerType.MATRIX_MULTIPLY or plugin_type == 'Gemm': |
| | fillcolor = 'orange' |
| | if fillcolor is not None: |
| | style = {**style, 'fillcolor': fillcolor} |
| | if no_style: |
| | style = {} |
| |
|
| | layer_attrs = {} |
| | layer_type = layer.type |
| | layer.to_subclass() |
| | if layer_type == trt.LayerType.CONSTANT: |
| | if not layer.is_shape_io: |
| | if trt.volume(layer.get_output(0).shape) <= 8: |
| | weights = layer.as_trt().weights |
| | if isinstance(weights, trt.Weights): |
| | weights = weights.numpy() |
| | value = np.array2string( |
| | weights, |
| | formatter={'float_kind': lambda x: f"{x:.2e}"}) |
| | layer_attrs['value'] = value |
| | elif layer_type == trt.LayerType.SHUFFLE: |
| | for attr_name in ['first_transpose', 'second_transpose']: |
| | attr_value = getattr(layer.as_trt(), attr_name) |
| | if tuple(attr_value) != (0, 1, 2, 3, 4, 5, 6, 7): |
| | tensor = layer.get_input( |
| | 0 |
| | ) if attr_name == 'first_transpose' else layer.get_output( |
| | 0) |
| | layer_attrs[attr_name] = tuple( |
| | attr_value)[:len(tensor.shape)] |
| | if layer.num_inputs < 2: |
| | attr_value = layer.as_trt().reshape_dims |
| | layer_attrs['reshape_dims'] = attr_value |
| | elif layer_type == trt.LayerType.SLICE: |
| | if layer.num_inputs < 2 or layer.get_input(1) is None: |
| | layer_attrs['start'] = layer.as_trt().start |
| | if layer.num_inputs < 4 or layer.get_input(3) is None: |
| | attr_value = layer.as_trt().stride |
| | if attr_value != tuple( |
| | [1] * len(layer.get_output(0).shape)): |
| | layer_attrs['stride'] = attr_value |
| | layer.to_base_class() |
| |
|
| | if layer.is_shape_io: |
| | labels = [layer.type.name] |
| | else: |
| | labels = [name, layer.type.name] |
| | for key, value in layer_attrs.items(): |
| | labels.append(f"{key}={value}") |
| | for attr_name in extra_attrs: |
| | if attr_name in layer.attrs: |
| | labels.append(f"{attr_name}={layer.attrs[attr_name]}") |
| | if layer.num_outputs == 1: |
| | output = layer.get_output(0) |
| | if output.name != f'{layer.name}_output_0': |
| | labels.append(f"output={output.name}") |
| | labels.extend(get_tensor_labels(output)) |
| | content = "\n".join(labels) |
| |
|
| | current_graph = get_device_graph(name) or graph |
| | current_graph = get_block_graph(layer, current_graph) |
| | current_graph.node(name, content, **style) |
| | layer_names.add(name) |
| |
|
| | for index, input in enumerate(layer.inputs): |
| | if input is not None: |
| | if input.is_graph_input or input.producer.num_outputs > 1: |
| | if input.name in tensor_names: |
| | graph.edge(input.name, name, str(index)) |
| | else: |
| | if input.producer.name in layer_names: |
| | graph.edge(input.producer.name, name, str(index)) |
| | if layer.num_outputs > 1 or (layer.num_outputs == 1 and |
| | layer.get_output(0).is_graph_output): |
| | for index, output in enumerate(layer.outputs): |
| | graph.edge(name, output.name, str(index)) |
| |
|
| | graph.subgraph(input_block_graph) |
| | graph.subgraph(output_block_graph) |
| | for parent_graph, block_graph in block_graph_mapping: |
| | parent_graph.subgraph(block_graph) |
| | for device_graph in device_graphs.values(): |
| | graph.subgraph(device_graph) |
| |
|
| | if not path: |
| | return graph.source |
| | graph.save(path) |
| |
|
| | @staticmethod |
| | def trt_to_dot(trt_network, path=None): |
| | graph = PipelineGraph.from_trt(trt_network) |
| | graph.assign_shapes() |
| | dot = graph.to_dot(no_style=True) |
| | if path is not None: |
| | with open(path, "w") as f: |
| | f.write(dot) |
| | else: |
| | return dot |
| |
|
