tensorrt 读取32位onnx
时间: 2023-08-24 13:04:51 浏览: 39
TensorRT支持读取32位的ONNX模型。在使用TensorRT API将ONNX模型转换为TensorRT引擎时,可以指定输入和输出数据类型。例如,可以使用以下代码将输入和输出数据类型设置为FP32:
```python
import tensorrt as trt
# Load ONNX model
onnx_model_file = "model.onnx"
with open(onnx_model_file, 'rb') as model_file:
onnx_model = model_file.read()
# Create TensorRT builder and network
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network()
# Set input and output data types to FP32
input_dtype = trt.float32
output_dtype = trt.float32
builder.set_fp16_mode(False)
builder.set_fp16_mode(False)
# Parse ONNX model and create TensorRT engine
parser = trt.OnnxParser(network, builder.logger)
parser.parse(onnx_model)
engine = builder.build_cuda_engine(network)
```
在上面的代码中,`builder.set_fp16_mode(False)`和`builder.set_int8_mode(False)`将TensorRT引擎的数据类型设置为FP32。如果想将数据类型设置为FP16或INT8,可以将这些方法的参数值设置为True。