colab中搭建tf-trt环境——解决TensorRT版本匹配问题

• Ubuntu18.04
• cuda11.1
• cudnn8
• python3.7
• tensorflow-gpu2.7.0
• TensorRT7.2.2
• pycuda2021.1

查看版本命令：

!lsb_release -a
"
No LSB modules are available.
Distributor ID:	Ubuntu
Description:	Ubuntu 18.04.5 LTS
Release:	18.04
Codename:	bionic
"

import tensorflow as tf
print(tf.__version__)
tf.sysconfig.get_build_info()
"
OrderedDict([('cpu_compiler', '/usr/bin/x86_64-linux-gnu-gcc-7'),
             ('cuda_compute_capabilities',
              ['compute_37',
               'compute_60',
               'compute_61',
               'compute_70',
               'compute_75',
               'compute_80']),
             ('cuda_version', '11.1'),
             ('cudnn_version', '8'),
             ('is_cuda_build', True),
             ('is_rocm_build', False),
             ('is_tensorrt_build', False)])
"

!python --version
"
Python 3.7.13
"

TensorRT安装包下载地址

• https://developer.nvidia.com/nvidia-tensorrt-download
• deb文件上传到云端磁盘
• tar文件下载到本地并解压，上传python文件夹中对应版本的whl文件

安装及测试

# 因模型使用2.8.0训练，低版本tf不支持模型中的部分操作
!pip install tensorflow-gpu==2.7.0# 自动删除以前版本

# 解压tensorrt-deb
!sudo dpkg -i './drive/MyDrive/Colab Notebooks/nv-tensorrt-repo-ubuntu1804-cuda11.1-trt7.2.2.3-ga-20201211_1-1_amd64.deb'

!sudo apt-get update

!sudo apt-get install libnvinfer7=7.2.2-1+cuda11.1 libnvinfer-plugin7=7.2.2-1+cuda11.1 libnvparsers7=7.2.2-1+cuda11.1  libnvonnxparsers7=7.2.2-1+cuda11.1 libnvinfer-bin=7.2.2-1+cuda11.1 libnvinfer-dev=7.2.2-1+cuda11.1 libnvinfer-plugin-dev=7.2.2-1+cuda11.1 libnvparsers-dev=7.2.2-1+cuda11.1 libnvonnxparsers-dev=7.2.2-1+cuda11.1 libnvinfer-samples=7.2.2-1+cuda11.1 libnvinfer-doc=7.2.2-1+cuda11.1

!sudo apt-mark hold libnvinfer7 libnvinfer-plugin7 libnvparsers7  libnvonnxparsers7 libnvinfer-bin libnvinfer-dev libnvinfer-plugin-dev libnvparsers-dev libnvonnxparsers-dev libnvinfer-samples libnvinfer-doc

!sudo apt-get install tensorrt

!pip install pycuda# 默认安装最新版本

# TensorRT.tar安装包中的python->.whl
!pip install './drive/MyDrive/Colab Notebooks/tensorrt-7.2.2.3-cp37-none-linux_x86_64.whl'

#检测tensorflow是否连接到对应tensorRT版本
import tensorflow as tf
print("TensorFlow version:", tf.__version__)
print("is_gpu_available:", tf.test.is_gpu_available())
import tensorflow.compiler as tf_cc
print("is TensorRT enabled:", tf_cc.tf2tensorrt._pywrap_py_utils.is_tensorrt_enabled())
print("loaded trt ver:", tf_cc.tf2tensorrt._pywrap_py_utils.get_loaded_tensorrt_version())
print("linked trt ver:", tf_cc.tf2tensorrt._pywrap_py_utils.get_linked_tensorrt_version())
'
TensorFlow version: 2.7.0
WARNING:tensorflow:From :3: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.list_physical_devices('GPU')` instead.
is_gpu_available: True
is TensorRT enabled: True
loaded trt ver: (7, 2, 2)
linked trt ver: (7, 2, 2)
'