DeepDeploy-Toolkit/test.py at master · xiaodiaoke/DeepDeploy-Toolkit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from utils.utils import process_image

from ONNX.python.onnx_run import ONNXModel
from TensorRT.python.trt_run import TRTModel
from RKNN.python.rknn_run import RKNNMolde
from Mnn.python.mnn_run import MNNModel


if __name__ == '__main__':
    img_path = "0510.jpg"
    img = process_image(img_path, (224, 224))

    # model = ONNXModel("net.onnx", device="gpu")
    # output = model.forward([img])
    # print(output[0].shape)

    # model = TRTModel("net.onnx", "net.plan", "fp16")
    # output = model.forward([img])
    # print(output[0].shape)

    model = TRTModel("net.onnx", "net.plan", "int8",
                     calibration_dataset="dataset.txt",
                     calibration_image_size=(224, 224),
                     calibration_mean=[0, 0, 0],
                     calibration_std=[1, 1, 1])
    output = model.forward([img])
    print(output[0].shape)

    # model = RKNNMolde("net.onnx", "net.rknn", "hybrid", quant_data_file="dataset.txt")
    # output = model.forward([img])
    # print(output[0].shape)
    # model.perf([img])

    # model = MNNModel("net.mnn")
    # output = model.forward([img])
    # print(output[0].shape)