没有合适的资源?快使用搜索试试~ 我知道了~
Pytorch 转TensorRT5
7 下载量 199 浏览量
2020-12-21
19:48:57
上传
评论
收藏 39KB PDF 举报
温馨提示
试读
3页
知乎上,说的也不错: https://zhuanlan.zhihu.com/p/88318324 https://blog.csdn.net/qq_38003892/article/details/89314108 1.torch2trt https://www.ctolib.com/amp/NVIDIA-AI-IOT-torch2trt.html https://github.com/DocF/YOLOv3-Torch2TRT https://github.com/traveller59/torch2trt 像官方代码: https://github.com/NVIDIA-AI-IOT/to
资源推荐
资源详情
资源评论
Pytorch 转转TensorRT5
知乎上,说的也不错:
https://zhuanlan.zhihu.com/p/88318324
https://blog.csdn.net/qq_38003892/article/details/89314108
1.torch2trt
https://www.ctolib.com/amp/NVIDIA-AI-IOT-torch2trt.html
https://github.com/DocF/YOLOv3-Torch2TRT
https://github.com/traveller59/torch2trt
像官方代码:
https://github.com/NVIDIA-AI-IOT/torch2trt
import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet
# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])
Execute
We can execute the returned TRTModule just like the original PyTorch model
y = model(x)
y_trt = model_trt(x)
# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))
Save and load
We can save the model as a state_dict.
torch.save(model_trt.state_dict(), 'alexnet_trt.pth')
We can load the saved model into a TRTModule
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))
Models
We tested the converter against these models using the test.sh script. You can generate the results by calling
2.Pytorch通过保存为ONNX模型转TensorRT5
这是另一篇博客:有开源代码
https://www.cnblogs.com/darkknightzh/p/11332155.html
https://github.com/darkknightzh/TensorRT_pytorch
1 Pytorch以ONNX方式保存模型
资源评论
weixin_38743084
- 粉丝: 12
- 资源: 931
上传资源 快速赚钱
- 我的内容管理 展开
- 我的资源 快来上传第一个资源
- 我的收益 登录查看自己的收益
- 我的积分 登录查看自己的积分
- 我的C币 登录后查看C币余额
- 我的收藏
- 我的下载
- 下载帮助
安全验证
文档复制为VIP权益,开通VIP直接复制
信息提交成功