Torch-TensorRT在Windows下编译安装的技巧

随笔7个月前发布 国主财经
101 0 0

Torch-TensorRT是pytorch官方对tensorrt的支持,将tensorrt推理加速功能直接集成进了pytorch。但是官方只提供linux预编译二进制,而缺乏对Windows编译支持,如果按照官方指引,不出意外将会得到一堆编译错误,本文记录一下几点技巧

1. 安装TensorRT,CUDA,CUDNN

本步骤略过,在Nvidia官网安装好之后添加到PATH即可

2. 指定TORCH_DIR

下载大于2.0版本的libtorch

Torch-TensorRT在Windows下编译安装的技巧

在cmake-gui中指向下载的libtorch目录,注意必须为libtorch/share/cmake/Torch

Torch-TensorRT在Windows下编译安装的技巧

修改mkl链接库

如果你没有安装mkl,可以直接使用conda install mkl安装,无需去intel官网注册下载,但要注意随后在Visual Studio中更改mkl链接库的名字。需要更改的有torch_trttorchtrt_plugintorchtrt_runtimetorchtrtc这几个解决方案,给mkl_*.lib加上_dll后缀即可。

Torch-TensorRT在Windows下编译安装的技巧

官方的Bug

官方的cmake其实留了一个bug,那就是core/lowering/passes忘了编译unpack_scaled_dot_product_attention.cpp,会导致后面找不到函数的reference。我已经修复了这个问题并给官方提交了PR,相信不久会合入master.

© 版权声明

相关文章

暂无评论

您必须登录才能参与评论!
立即登录
暂无评论...