TensorRT quickstart

·  阅读 255

docs.nvidia.com/deeplearnin…

github.com/NVIDIA/Tens…

  • 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

在 C++ 中运行engine

分类:
开发工具
标签:
分类:
开发工具
标签:
收藏成功!
已添加到「」, 点击更改