### TensorRT 中获取 Engine 维度的函数 在 TensorRT 中，获取 Engine 的维度信息可以通过特定的 API 实现。这些 API 主要集中在 `ICudaEngine` 和 `IExecutionContext` 接口中[^1]。以下是一些常用的函数及其功能描述： - **`getBindingDimensions(int bindingIndex)`**：此函数用于获取指定绑定索引（binding index）的张量维度。返回值是一个 `Dims` 类型的对象，表示该绑定的维度信息[^2]。 - **`getProfileDimensions(int bindingIndex, int profileIndex, OptimizationProfileSelector selector)`**：如果使用了优化配置文件（Optimization Profiles），可以通过此函数获取不同优化配置下的最小、最大或优化维度。参数 `selector` 可以是 `OPTIMIZATION_PROFILE_MIN`, `OPTIMIZATION_PROFILE_MAX` 或 `OPTIMIZATION_PROFILE_OPT`[^3]。 - **`getNbBindings()`**：用于获取绑定的数量。结合 `getBindingDimensions` 函数，可以遍历所有绑定并获取其维度信息[^2]。 - **`getDeviceMemorySize()`**：虽然不直接与维度相关，但此函数可以获取引擎所需的设备内存大小，间接反映模型的复杂度和规模。 以下是通过 C++ 使用上述 API 获取 Engine 绑定维度的示例代码： ```cpp #include <NvInfer.h> #include <iostream> void printBindingDimensions(nvinfer1::ICudaEngine& engine) { int nbBindings = engine.getNbBindings(); // 获取绑定数量 for (int i = 0; i < nbBindings; ++i) { nvinfer1::Dims dims = engine.getBindingDimensions(i); // 获取绑定维度 std::cout << "Binding " << i << " dimensions: "; for (int j = 0; j < dims.nbDims; ++j) { std::cout << dims.d[j] << " "; } std::cout << std::endl; } } ``` ### 注意事项 在使用这些 API 时需要注意以下几点： - 绑定索引（binding index）对应输入和输出张量，通常按照定义顺序排列。 - 如果使用动态形状（Dynamic Shapes），需要确保正确配置优化配置文件，并使用 `getProfileDimensions` 获取动态范围内的维度信息[^3]。