TensorRT是可以在NVIDIA各种GPU硬件平台下运行的一个C++推理框架。我们利用Pytorch、TF或者其他框架训练好的模型,可以转化为TensorRT的格式,然后利用TensorRT推理引擎去运行我们这个模型,从而提升这个模型在英伟达GPU上运行的速度。速度提升的比例是比较可观的。
对于GPU来说,因为底层的硬件设计,更适合并行计算也更喜欢密集型计算。TensorRT所做的优化也是基于GPU进行优化,当然也是更喜欢那种一大块一大块的矩阵运算,尽量直通到底。因此对于通道数比较多的卷积层和反卷积层,优化力度是比较大的;如果是比较繁多复杂的各种细小op操作(例如reshape、gather、split等),那么TensorRT的优化力度就没有那么夸张了。
为了更充分利用GPU的优势,我们在设计模型的时候,可以更加偏向于模型的并行性,因为同样的计算量,“大而整”的GPU运算效率远超“小而碎”的运算。
左上角是原始网络(googlenet),右上角相对原始层进行了垂直优化,将conv+bias(BN)+relu进行了融合优化;而右下角进行了水平优化,将所有1x1的CBR融合成一个大的CBR;左下角则将concat层直接去掉,将contact层的输入直接送入下面的操作中,不用单独进行concat后在输入计算,相当于减少了一次传输吞吐。