推理加速1_pytorch转tensorRT

1.模型转换(pytorch模型->onnx->tensorRT)

1.1 pytorch->onnx

dummy_image = torch.randn(1, 3, 224, 224)
model.to(device)
dummy_image.to(device)
dest = "test.onnx"
torch.onnx.export(model, dummy_image, dest, export_params=True, input_names=['modelInput'],  output_names=['modelOutput'], opset_version=14)

# Now run onnxruntime to verify
import onnxruntime as ort

ori_out = model(dummy_image)
ori = ori_out[0].detach().numpy()

ort_sess = ort.InferenceSession(dest)
ort_inputs = {ort_sess.get_inputs()[0].name: dummy_image.numpy()}
onnx = ort_sess.run(None, ort_inputs)[0].reshape(-1)

print("ori=", ori)
print("onxx=", onnx)

# 2、onnx to tensorTT
https://zhuanlan.zhihu.com/p/395590559