forked from jiuyuan/InfiniTensor
30 lines
965 B
Python
30 lines
965 B
Python
import sys
|
|
import onnx
|
|
import torch
|
|
import numpy as np
|
|
from pyinfinitensor.onnx import OnnxStub, backend
|
|
|
|
if __name__ == '__main__':
|
|
args = sys.argv
|
|
if len(sys.argv) != 2:
|
|
print("Usage: python onnx_inference.py model_name.onnx")
|
|
exit()
|
|
model_path = sys.argv[1]
|
|
# print(model_path)
|
|
|
|
onnx_model = onnx.load(model_path)
|
|
onnx_input = onnx_model.graph.input[0]
|
|
input_shape = [[d.dim_value for d in _input.type.tensor_type.shape.dim]
|
|
for _input in onnx_model.graph.input]
|
|
# Assume that there is only one input tensor
|
|
input_shape = input_shape[0]
|
|
# print(input_shape)
|
|
input_data = np.random.random(input_shape).astype(np.float32)
|
|
|
|
model = OnnxStub(onnx_model, backend.cuda_runtime())
|
|
next(iter(model.inputs.values())).copyin_numpy(input_data)
|
|
model.run()
|
|
outputs = next(iter(model.outputs.values())).copyout_numpy()
|
|
outputs = torch.tensor(outputs)
|
|
print(outputs.shape)
|