- generate the TensorRT engine from a PyTorch model exported to ONNX and converted using
trtexec - C++ runtime API – run inference using engine and TensorRT’s C++ API
- Python runtime API – run inference using engine and TensorRT’s Python API
从pytorch模型导出trt engine
- pytorch to onnx
# FC-ResNet101 pretrained model from torch-hub extended with argmax layer
class FCN_ResNet101(nn.Module):
def __init__(self):
super(FCN_ResNet101, self).__init__()
self.model = torch.hub.load('pytorch/vision:v0.6.0', 'fcn_resnet101', pretrained=True)
def forward(self, inputs):
x = self.model(inputs)['out']
x = x.argmax(1, keepdims=True)
return x
model = FCN_ResNet101()
model.eval()
# Generate input tensor with random values
input_tensor = torch.rand(4, 3, 224, 224)
# Export torch model to ONNX
print("Exporting ONNX model {}".format(output_onnx))
torch.onnx.export(model, input_tensor, output_onnx,
opset_version=12,
do_constant_folding=True,
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: "batch", 2: "height", 3: "width"},
"output": {0: "batch", 2: "height", 3: "width"}},
verbose=False)
附加一个 argmax 输出层以生成最高概率的每种像素类标签
- onnx to trt engine
trtexec --onnx=fcn-resnet101.onnx --explicitBatch --fp16 --workspace=64 \
--minShapes=input:1x3x256x256 --optShapes=input:1x3x1026x1282 --maxShapes=input:1x3x1440x2560 \
--buildOnly --saveEngine=fcn-resnet101.engine
- 随机输入验证模型可运行
trtexec --shapes=input:1x3x1026x1282 --loadEngine=fcn-resnet101.engine
