pytorch转onnx, onnx转tensorrt

从pytorch转换到onnx - 知乎

pytorch模型转onnx模型_挣扎的笨鸟的博客-CSDN博客_pytorch转onnx

先看结果：

pytorch、onnx、tensorrt模型输出完全一样，done

 

import torch
import torchvision.transforms as transforms
from PIL import Image
import onnxruntime
from model import LeNet


import torch
import torch.nn
import onnx
import onnx
import onnxruntime as ort
import numpy as np

def main():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    net = LeNet()
    net.load_state_dict(torch.load('Lenet.pth'))

    net.eval()

    input_names = ['input']
    output_names = ['output']

    x = torch.randn(1, 3, 32, 32, requires_grad=True)

    torch.onnx.export(net, x, 'best.onnx', input_names=input_names, output_names=output_names, verbose='True')


    im = Image.open('1.jpg')
    im = transform(im)  # [C, H, W]
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]

    img1 = np.array(im, dtype=np.float32)
    model = onnx.load('best.onnx')
    onnx.checker.check_model(model)

    # session = ort.InferenceSession('best.onnx')
    x = np.random.randn(1, 3, 32, 32).astype(np.float32)  # 注意输入type一定要np.float32!!!!!
    # x= torch.randn(batch_size,chancel,h,w)

    # session = onnxruntime.InferenceSession('best.onnx')
    session = onnxruntime.InferenceSession('best.onnx', providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider',
                                                       'CPUExecutionProvider'])

    # session = onnx.load(onnx_path)
    print("The model expects input shape: ", session.get_inputs()[0].shape)

    outputs = session.run(None, input_feed={'input': img1})
    print(11111111111, outputs)


    with torch.no_grad():
        outputs = net(im)
        print(222222222222222, outputs)
        # predict = torch.max(outputs, dim=1)[1].numpy()
    # print(classes[int(predict)])


if __name__ == '__main__':
    main()

#
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import sys
import torch
import torchvision.transforms as transforms
from PIL import Image
import onnxruntime
from model import LeNet


import torch
import torch.nn
import onnx
import onnx
import onnxruntime as ort
import numpy as np
# This sample uses an MNIST PyTorch model to create a TensorRT Inference Engine
import model
import numpy as np
import pycuda.autoinit
import tensorrt as trt

sys.path.insert(1, os.path.join(sys.path[0], ".."))
import common

# You can set the logger severity higher to suppress messages (or lower to display more messages).
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)


class ModelData(object):
    INPUT_NAME = "data"
    INPUT_SHAPE = (1, 1, 28, 28)
    OUTPUT_NAME = "prob"
    OUTPUT_SIZE = 10
    DTYPE = trt.float32


def populate_network(network, weights):
    # Configure the network layers based on the weights provided.
    input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)

    def add_matmul_as_fc(net, input, outputs, w, b):
        assert len(input.shape) >= 3
        m = 1 if len(input.shape) == 3 else input.shape[0]
        k = int(np.prod(input.shape) / m)
        assert np.prod(input.shape) == m * k
        n = int(w.size / k)
        assert w.size == n * k
        assert b.size == n

        input_reshape = net.add_shuffle(input)
        input_reshape.reshape_dims = trt.Dims2(m, k)

        filter_const = net.add_constant(trt.Dims2(n, k), w)
        mm = net.add_matrix_multiply(
            input_reshape.get_output(0),
            trt.MatrixOperation.NONE,
            filter_const.get_output(0),
            trt.MatrixOperation.TRANSPOSE,
        )

        bias_const = net.add_constant(trt.Dims2(1, n), b)
        bias_add = net.add_elementwise(mm.get_output(0), bias_const.get_output(0), trt.ElementWiseOperation.SUM)

        output_reshape = net.add_shuffle(bias_add.get_output(0))
        output_reshape.reshape_dims = trt.Dims4(m, n, 1, 1)
        return output_reshape

    conv1_w = weights["conv1.weight"].numpy()
    conv1_b = weights["conv1.bias"].numpy()
    conv1 = network.add_convolution(
        input=input_tensor, num_output_maps=20, kernel_shape=(5, 5), kernel=conv1_w, bias=conv1_b
    )
    conv1.stride = (1, 1)

    pool1 = network.add_pooling(input=conv1.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), 50, (5, 5), conv2_w, conv2_b)
    conv2.stride = (1, 1)

    pool2 = network.add_pooling(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
    pool2.stride = (2, 2)

    fc1_w = weights["fc1.weight"].numpy()
    fc1_b = weights["fc1.bias"].numpy()
    fc1 = add_matmul_as_fc(network, pool2.get_output(0), 500, fc1_w, fc1_b)

    relu1 = network.add_activation(input=fc1.get_output(0), type=trt.ActivationType.RELU)

    fc2_w = weights["fc2.weight"].numpy()
    fc2_b = weights["fc2.bias"].numpy()
    fc2 = add_matmul_as_fc(network, relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b)

    fc2.get_output(0).name = ModelData.OUTPUT_NAME
    network.mark_output(tensor=fc2.get_output(0))


def build_engine(img):
    # For more information on TRT basics, refer to the introductory samples.
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(common.EXPLICIT_BATCH)
    # config = builder.create_builder_config()

    parser = trt.OnnxParser(network, TRT_LOGGER)
    builder.max_workspace_size = common.GiB(1)

    with open('best.onnx', 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: Failed to parse the ONNX file.')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
    engine = builder.build_cuda_engine(network)
    context = engine.create_execution_context()
    inputs, outputs, bindings, stream = common.allocate_buffers(engine)

    #img1 = np.asarray(img).transpose([0, 3, 1, 2]).astype(trt.nptype(trt.float32)).ravel()
    img1 = img.ravel()
    np.copyto(inputs[0].host, img1)

    with open('best.engine', 'wb') as t:
        t.write(engine.serialize())
    trt_outputs = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs,
                                         stream=stream)
    print(trt_outputs)

    # runtime = trt.Runtime(TRT_LOGGER)
    #
    # config.max_workspace_size = common.GiB(1)
    # # Populate the network using weights from the PyTorch model.
    # populate_network(network, weights)
    # # Build and return an engine.
    # plan = builder.build_serialized_network(network, config)
    # return runtime.deserialize_cuda_engine(plan)


# Loads a random test case from pytorch's DataLoader
def load_random_test_case(model, pagelocked_buffer):
    # Select an image at random to be the test case.
    img, expected_output = model.get_random_testcase()
    # Copy to the pagelocked input buffer
    np.copyto(pagelocked_buffer, img)
    return expected_output



def main():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)),
         transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    net = LeNet()
    net.load_state_dict(torch.load('Lenet.pth'))

    net.eval()

    input_names = ['input']
    output_names = ['output']

    x = torch.randn(1, 3, 32, 32, requires_grad=True)

    torch.onnx.export(net, x, 'best.onnx', input_names=input_names, output_names=output_names, verbose='True')


    im = Image.open('1.jpg')
    im = transform(im)  # [C, H, W]
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]

    img1 = np.array(im, dtype=np.float32)
    model = onnx.load('best.onnx')
    onnx.checker.check_model(model)
    # print(2211122222222222222, onnx.helper.printable_graph(model.graph))

    # session = ort.InferenceSession('best.onnx')
    x = np.random.randn(1, 3, 32, 32).astype(np.float32)  # 注意输入type一定要np.float32!!!!!
    # x= torch.randn(batch_size,chancel,h,w)

    # session = onnxruntime.InferenceSession('best.onnx')
    session = onnxruntime.InferenceSession('best.onnx', providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider',
                                                       'CPUExecutionProvider'])

    # session = onnx.load(onnx_path)
    # print("The model expects input shape: ", session.get_inputs()[0].shape)

    outputs = session.run(None, input_feed={'input': img1})
    print(11111111111, outputs)


    with torch.no_grad():
        outputs = net(im)
        print(222222222222222, outputs)
        # predict = torch.max(outputs, dim=1)[1].numpy()
    # print(classes[int(predict)])

    build_engine(img1)

if __name__ == "__main__":
    main()