Layer is a wrapper for TensorRT's ILayer with several python-friendly helper functions.
| 13 | |
| 14 | |
| 15 | class Layer: |
| 16 | """Layer is a wrapper for TensorRT's ILayer with several python-friendly helper functions.""" |
| 17 | |
| 18 | def __init__(self, network: Network, trt_layer: trt.ILayer): |
| 19 | self._network = weakref.ref(network) |
| 20 | self.trt_layer = trt_layer |
| 21 | |
| 22 | assert isinstance(self.network, Network) |
| 23 | assert isinstance(self.trt_layer, trt.ILayer) |
| 24 | |
| 25 | @property |
| 26 | def network(self): |
| 27 | return self._network() |
| 28 | |
| 29 | def get_inputs(self, *indices: int): |
| 30 | """Get the input tensors of the layer. |
| 31 | |
| 32 | Parameters: |
| 33 | idx: the indices of the input tensor, will return all inputs if left empty |
| 34 | |
| 35 | Returns: |
| 36 | List[Tensor] |
| 37 | """ |
| 38 | from .functional import Tensor |
| 39 | |
| 40 | indices = indices if indices else range(self.trt_layer.num_inputs) |
| 41 | |
| 42 | ret = [] |
| 43 | for i in indices: |
| 44 | assert i < self.trt_layer.num_inputs, ( |
| 45 | f"Invalid input index {i} for layer {self.trt_layer.name}" |
| 46 | ) |
| 47 | |
| 48 | tensor = self.trt_layer.get_input(i) |
| 49 | tensor = Tensor(trt_tensor=tensor, network=self.network, is_network_input=False) |
| 50 | ret.append(tensor) |
| 51 | return ret |
| 52 | |
| 53 | def get_outputs(self, *indices: int): |
| 54 | """Get the output tensor of the layer. |
| 55 | |
| 56 | Parameters: |
| 57 | idx: the index of the output tensor |
| 58 | |
| 59 | Returns: |
| 60 | List[Tensor] |
| 61 | """ |
| 62 | from .functional import Tensor |
| 63 | |
| 64 | indices = indices if indices else range(self.trt_layer.num_outputs) |
| 65 | |
| 66 | ret = [] |
| 67 | for i in indices: |
| 68 | assert i < self.trt_layer.num_outputs, ( |
| 69 | f"Invalid output index {i} for layer {self.trt_layer.name}" |
| 70 | ) |
| 71 | |
| 72 | tensor = self.trt_layer.get_output(i) |
no test coverage detected