Search code examples
python-3.xtensorflowtf.kerastensorrt

convert a keras sequential model to tensor rt


How to convert a keras sequential model to tensor-rt? I am working on face recognition, I have created model using keras sequential model on facial encoding. Now I want to implement that on jtson xavier. This is the model architecture:

    def Simple_NeuralNet(number_classes):
        model = Sequential()
        model.add(Dense(units=256, activation='relu', input_dim=128))
        model.add(BatchNormalization())
        model.add(Dense(units=1024, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(units=512, activation='relu'))
        model.add(BatchNormalization())
        model.add(Dense(units=number_classes, activation='softmax'))
        return model

Solution

  • Checkout tf2onnx (https://github.com/onnx/tensorflow-onnx) or keras2onnx (https://github.com/onnx/keras-onnx) to convert your model to ONNX format. There are more some details on conversion to ONNX in this post: https://stackoverflow.com/a/59793636/10993413.

    Then you should be able to quickly verify your model using:

    trtexec --onnx=model.onnx
    

    or the TensorRT API to convert your model in TensorRT 6.0 (Jetson).

    If using TensorRT 7.0, be sure to add the --explicitBatch flag to trtexec, or NetworkDefinitionCreationFlag.EXPLICIT_BATCH if using the API:

    trtexec --explicitBatch --onnx=model.onnx
    

    The TensorRT release comes with some samples for both C++ and Python that demonstate how to convert an ONNX model to TensorRT using the API. The samples can be found in TRT_RELEASE_DIR/samples.