Method _set_layer_name

tensorrt_llm/network.py:253–299

(self, layer)

        return self._named_parameters

    def _set_layer_name(self, layer):
        original_layer_name = layer.name
        layer_name = str(layer.type).split('.')[-1]
        current_module = self._module_call_stack.get_current_module()

        func_stack = []
        frame = inspect.currentframe().f_back.f_back
        while frame:
            func_name = frame.f_code.co_name
            line_num = frame.f_lineno
            if func_name == "forward":
                break
            func_stack.insert(0, f"{func_name}_L{line_num}")
            if len(func_stack) >= 10:
                # NOTE: TRT error messages has a character limit.
                #       Limiting to only 10 levels helps retain
                #       the true error message from TRT.
                break
            frame = frame.f_back
        current_module = f"{current_module}.{'.'.join(func_stack)}"

        if layer.type == trt.LayerType.PLUGIN_V2:
            layer_name = '_'.join(
                [layer_name,
                 str(layer.plugin.plugin_type).split('.')[-1]])
        elif layer.type in [
                trt.LayerType.UNARY, trt.LayerType.REDUCE,
                trt.LayerType.ELEMENTWISE
        ]:
            layer_name = '_'.join([layer_name, str(layer.op).split('.')[-1]])

        layer.name = self._name_generator(layer_name, current_module)
        for idx in range(layer.num_outputs):
            # TRT initializes tensor names from the initial layer's name when the layer is created,
            # and does not update tensor names when layer name changed by application, needs to
            # change the tensor name to align with the new layer name for better debugging
            layer.get_output(idx).name = f"{layer.name}_output_{idx}"
        if original_layer_name != layer.name:
            if layer.type == trt.LayerType.PLUGIN_V2:
                plugin_info = get_plugin_info(self.trt_network,
                                              original_layer_name)
                if plugin_info is not None:
                    set_plugin_info(self.trt_network, layer.name, plugin_info)
                    delete_plugin_info(self.trt_network, original_layer_name)

        # Set layer metadata to the same as the layer name so that it can show up in NVTX.
        layer.metadata = layer.name

    def register_ndarray(self, ndarray: np.ndarray) -> None:
        ''' When the functional APIs need to create local numpy array and use as weights for constant or other layers,