推理加速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

__EOF__

本文作者：userName
本文链接：https://www.cnblogs.com/pyclq/p/17393374.html
关于博主：评论和私信会在第一时间回复。或者直接私信我。
版权声明：本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
声援博主：如果您觉得文章对您有帮助，可以点击文章右下角【推荐】一下。您的鼓励是博主的最大动力！