Creating engines for PyTorch or onnx
mlcoop opened this issue · 8 comments
I have models built in PyTorch and also converted into onnx.
PyTorch 1.0
TensortRt 51.5.0
First, tried to build engine as it was built in mnist example.
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
class ModelDataWatermark(object): INPUT_NAME = "data" INPUT_SHAPE = (1, 128, 128) OUTPUT_NAME = "prob" OUTPUT_SIZE = 2 DTYPE = trt.float32
`
def populate_network_watermark(network, weights):
input_tensor = network.add_input(name=ModelDataWatermark.INPUT_NAME,
dtype=ModelDataWatermark.DTYPE, shape=ModelDataWatermark.INPUT_SHAPE)
conv1_w = weights['conv1.weight'].numpy()
conv1_b = weights['conv1.bias'].numpy()
conv1 = network.add_convolution(input=input_tensor, num_output_maps=64, kernel_shape=(3, 3), kernel=conv1_w, bias=conv1_b)
conv1.stride = (1, 1)
relu_1 = network.add_activation(input=conv1.get_output(0), type=trt.ActivationType.RELU)
pool1 = network.add_pooling(input=relu_1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2))
pool1.stride = (2, 2)
conv2_w = weights['conv2.weight'].numpy()
conv2_b = weights['conv2.bias'].numpy()
conv2 = network.add_convolution(pool1.get_output(0), num_output_maps=128, kernel_shape=(3, 3), kernel=conv2_w, bias=conv2_b)
conv2.stride = (1, 1)
relu_2 = network.add_activation(input=conv2.get_output(0), type=trt.ActivationType.RELU)
pool2 = network.add_pooling(relu_2.get_output(0), trt.PoolingType.MAX, (2, 2))
pool2.stride = (2, 2)
conv3_w = weights['conv3.weight'].numpy()
conv3_b = weights['conv3.bias'].numpy()
conv3 = network.add_convolution(pool2.get_output(0), 256, (3, 3), conv3_w, conv3_b)
conv3.stride = (1, 1)
relu_3 = network.add_activation(input=conv3.get_output(0), type=trt.ActivationType.RELU)
pool3 = network.add_pooling(relu_3.get_output(0), trt.PoolingType.MAX, (2, 2))
pool3.stride = (2, 2)
conv4_w = weights['conv4.weight'].numpy()
conv4_b = weights['conv4.bias'].numpy()
conv4 = network.add_convolution(pool3.get_output(0), 512, (3, 3), conv4_w, conv4_b)
conv4.stride = (1, 1)
relu_4 = network.add_activation(input=conv4.get_output(0), type=trt.ActivationType.RELU)
pool4 = network.add_pooling(relu_4.get_output(0), trt.PoolingType.MAX, (2, 2))
pool4.stride = (2, 2)
conv5_w = weights['conv5.weight'].numpy()
conv5_b = weights['conv5.bias'].numpy()
conv5 = network.add_convolution(pool4.get_output(0), 1024, (3, 3), conv5_w, conv5_b)
conv5.stride = (1, 1)
relu_5 = network.add_activation(input=conv5.get_output(0), type=trt.ActivationType.RELU)
pool5 = network.add_pooling(relu_5.get_output(0), trt.PoolingType.MAX, (2, 2))
pool5.stride = (2, 2)
fc1_w = weights['fc1.weight'].numpy()
fc1_b = weights['fc1.bias'].numpy()
fc1 = network.add_fully_connected(input=pool5.get_output(0), num_outputs=4096, kernel=fc1_w, bias=fc1_b)
fc2_w = weights['fc2.weight'].numpy()
fc2_b = weights['fc2.bias'].numpy()
fc2 = network.add_fully_connected(input=fc1.get_output(0), num_outputs=4096, kernel=fc2_w, bias=fc2_b)
fc3_w = weights['fc3.weight'].numpy()
fc3_b = weights['fc3.bias'].numpy()
fc3 = network.add_fully_connected(input=fc2.get_output(0), num_outputs=ModelDataWatermark.OUTPUT_SIZE, kernel=fc3_w, bias=fc3_b)
fc3.get_output(0).name = ModelDataWatermark.OUTPUT_NAME
network.mark_output(tensor=fc3.get_output(0))`
checkpoints = '/home/malibayev/kolesa-image-moderation-sanic/static/models/torch_model_07_06.pt' model = WatermarkModel(num_classes=2) model.load_state_dict(torch.load(checkpoints)) weights = model.state_dict()
def build_engine(weights): # For more information on TRT basics, refer to the introductory samples. with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network: builder.max_workspace_size = common.GiB(5) # Populate the network using weights from the PyTorch model. populate_network_watermark(network, weights) # Build and return an engine. return builder.build_cuda_engine(networ
engine = build_engine(weights)
AttributeError: enter was given
Then I tried to use Onnx Parser.
class ModelData(object): MODEL_PATH = "/home/malibayev/kolesa-image-moderation-sanic/text_clf_opset7.onnx" INPUT_SHAPE = (3, 224, 224) # We can convert TensorRT data types to numpy types with trt.nptype() DTYPE = trt.float32
builder = trt.Builder(TRT_LOGGER) network = builder.create_network() parser = trt.OnnxParser(network, TRT_LOGGER) builder.max_workspace_size = common.GiB(1) model = open('/home/malibayev/kolesa-image-moderation-sanic/text_clf_opset7.onnx', 'rb') parser.parse(model.read()) engine = builder.build_cuda_engine(network)
which also returned None
parser.parse(model.read()) returns False
Please Help
Did you follow the PyTorch instructions to export to a .onnx file?
https://pytorch.org/docs/stable/onnx.html
It appears that you did as I see you may have generated text_clf_opset7.onnx
Once you have the .onnx model, you can use trtexec
found inside the NGC TensorRT container to build a TRT engine file. There is an example here: https://github.com/NVIDIA/tensorrt-laboratory/blob/master/examples/ONNX/resnet50/build.py
try this from the directory that contains your .onnx file, replacing name-of-your.onnx
ls *.onnx
nvidia-docker run --rm -ti -v $PWD:/work --workdir /work nvcr.io/nvidia/tensorrt:19.04-py3 bash
ls *.onnx
trtexec --onnx=name-of-your.onnx --batch=1 --saveEngine=test.engine
@ryanolson hey!
`root@375cb0eb2afa:/work# trtexec --onnx=watermark.onnx --batch=1 --saveEngine=test.engine
&&&& RUNNING TensorRT.trtexec # trtexec --onnx=watermark.onnx --batch=1 --saveEngine=test.engine
[I] onnx: watermark.onnx
[I] batch: 1
[I] saveEngine: test.engine
Input filename: watermark.onnx
ONNX IR version: 0.0.3
Opset version: 9
Producer name: pytorch
Producer version: 0.4
Domain:
Model version: 0
Doc string:
Unsupported ONNX data type: INT64 (7)
[E] [TRT] Parameter check failed at: ../builder/Layers.cpp::ConstantLayer::1585, condition: weights.type == DataType::kFLOAT || weights.type == DataType::kHALF || weights.type == DataType::kINT32
While parsing node number 17 [Gather]:
ERROR: onnx2trt_utils.hpp:275 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims
[E] failed to parse onnx file
[E] Engine could not be created
[E] Engine could not be created
&&&& FAILED TensorRT.trtexec # trtexec --onnx=watermark.onnx --batch=1 --saveEngine=test.engine`
Shoul I rewrite model in Pytorch 1.0 ?
Let's tackle the first error:
Unsupported ONNX data type: INT64 (7)
Is there anywhere in the model definition that you are explicitly using python int
or using np.int64
? If you have python list/array and convert it to numpy, it will be an np.int64
value.
>>> import numpy as np
>>> a = [1]
>>> b = np.asarray(a)
>>> b
array([1])
>>> b.dtype
dtype('int64')
Try to weed out any integer values that might be set to np.int64
since TensorRT only support INT32 and INT8.
@ryanolson ,
The model is pretty simple, without any numpy in it. Only nn.Module from PyTorch.
TensorRT doesn't support np.int64
, so we need to determine which bits are using np.int64
and convert them to np.int32
So I converted all layers into torch.FloatTensor then did
root@33aa966b729c:/work# trtexec --model='watermark_float.onnx' --batch=1 --saveEngine='test.engine' &&&& RUNNING TensorRT.trtexec # trtexec --model=watermark_float.onnx --batch=1 --saveEngine=test.engine [I] model: watermark_float.onnx [I] batch: 1 [I] saveEngine: test.engine [E] Deploy file not specified [E] Engine could not be created &&&& FAILED TensorRT.trtexec # trtexec --model=watermark_float.onnx --batch=1 --saveEngine=test.engine
It fails with "Deploy file not specified"
Change --model
to --onnx
Please open if you need more assistance.