ubuntu 安装 jax jaxlib cpu 和 gpu 版本 以及 tensorflow tensorRT的安装

需要事先安装较新版本的 cuda 和cudnn ，例如11.8 + 8.8

在已经安装过cuda的机器上安装新版cuda sdk 和 cudnn 可参考前述：

ubuntu 安装 多版本 cuda 11.4 11.8_eloudy的博客-csdn博客

一，安装python3 和 pip3

sudo apt-get install python3

sudo apt-get install python3-pip

pip3 --version

sudo pip3 install --upgrade pip

二，安装 cpu版本的 jax 和 jaxlib
  

  使用pip官方源安装 jax:  

sudo pip3 install jax jaxlib

    使用pip清华源安装 jax:(更新很及时，没看出版本差别)

sudo pip3 install jax jaxlib -i https://pypi.tuna.tsinghua.edu.cn/simple

        其他可选源：

             -i https://mirrors.aliyun.com/pypi/simple/
             -i http://mirrors.cloud.tencent.com/pypi/simple
             -i http://pypi.douban.com/simple

        永久换pip源：

sudo pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

        配置其他网址，类似替换
        
    安装内容：
        downloading jax-0.2.17.tar.gz
        downloading jaxlib-0.1.69-cp36-none-manylinux2010_x86_64.whl
        downloading absl_py-1.4.0-py3-none-any.whl
        downloading numpy-1.19.5-cp36-cp36m-manylinux2010_x86_64.whl
        downloading opt_einsum-3.3.0-py3-none-any.whl
        downloading flatbuffers-2.0.7-py2.py3-none-any.whl
        downloading scipy-1.5.4-cp36-cp36m-manylinux1_x86_64.whl

    成功安装后会显示：
successfully built jax
installing collected packages: scipy, opt_einsum, jaxlib, jax
successfully installed jax-0.4.4 jaxlib-0.4.4 opt_einsum-3.3.0 scipy-1.10.1
        

三，安装gpu版本的jax和jaxlib

安装gpu版本时遇到了一些版本问题，最后干脆将系统升级到cuda 11.8  cudnn 8.8

ubuntu 安装 多版本 cuda 11.4 11.8_eloudy的博客-csdn博客

先卸载原先的jax和 jaxlib，并更新pip：

pip uninstall jax jaxlib

pip install --upgrade pip

然后直接显式指定cuda和cudnn版本安装：

pip install "jax[cuda11_cudnn86]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

安装成功：

四，示例测试

1. 

python3
>>> import jax.numpy as np
>>>
 

2. 

import jax.numpy as np

from jax import random
import time

x = random.uniform(random.prngkey(0), [5000, 5000])
st=time.time()
try:
        y=np.matmul(x, x)

except exception:

    print("error")

print(time.time()-st)

print(y)

运行时没有向cpu版本那样提示找不到gpu：

3. 生成三维数组

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

key = random.prngkey(0)
x = random.normal(key,(7, 2, 3))
print(x)

4. 计算 lu分解

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

import pprint
import scipy
import scipy.linalg

key = random.prngkey(0)
a = random.normal(key,(7, 7))
print("a=")
pprint.pprint(a)

p, l, u = scipy.linalg.lu(a)

print("a=")
pprint.pprint(a)

print("p=")
pprint.pprint(p)

print("l=")
pprint.pprint(l)

print("u=")
pprint.pprint(u)

执行： 

官方安装参考网址

github - google/jax: composable transformations of python+numpy programs: differentiate, vectorize, jit to gpu/tpu, and more

五，顺便安装个tensorflow

先炸一炸最新版本

$ pip3 install tensorflow-gpu==33.33

输出如下：

mklguy@mklguy--pc:~/ex/tensorflow_ex$ pip install tensorflow-gpu==33.33
defaulting to user installation because normal site-packages is not writeable
error: could not find a version that satisfies the requirement tensorflow-gpu==33.33 (from versions: 2.2.0, 2.2.1, 2.2.2, 2.2.3, 2.3.0, 2.3.1, 2.3.2, 2.3.3, 2.3.4, 2.4.0, 2.4.1, 2.4.2, 2.4.3, 2.4.4, 2.5.0, 2.5.1, 2.5.2, 2.5.3, 2.6.0, 2.6.1, 2.6.2, 2.6.3, 2.6.4, 2.6.5, 2.7.0rc0, 2.7.0rc1, 2.7.0, 2.7.1, 2.7.2, 2.7.3, 2.7.4, 2.8.0rc0, 2.8.0rc1, 2.8.0, 2.8.1, 2.8.2, 2.8.3, 2.8.4, 2.9.0rc0, 2.9.0rc1, 2.9.0rc2, 2.9.0, 2.9.1, 2.9.2, 2.9.3, 2.10.0rc0, 2.10.0rc1, 2.10.0rc2, 2.10.0rc3, 2.10.0, 2.10.1, 2.11.0rc0, 2.11.0rc1, 2.11.0rc2, 2.11.0, 2.12.0)
error: no matching distribution found for tensorflow-gpu==33.33

可得当前最大版本为2.12.0；

选择tensorflow 2.11.0吧：

$ pip3 install tensorflow-gpu==2.11.0

 安装完成后的效果如图：

测试：

python3

>>> import tensorflow as tf

发现提示warning，需要安装tensorrt；

六，安装tensorrt

步骤：

1. 下载 8.6cudnn版的linux x64 的 tar包， tensorrt-8.5.3.1.linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz

在该网址下载 

https://developer.nvidia.com/nvidia-tensorrt-8x-download

2. 解压到对应cuda sdk版本的合适的位置

例如 /usr/local/cuda-11.8/

3. 设置环境变量

export path=/usr/local/cuda/bin:$path
export ld_library_path=/usr/local/cuda/lib64:$ld_library_path
export ld_library_path=$ld_library_path:/usr/local/cuda-11.8/tensorrt_sdk/tensorrt-8.5.3.1/lib
export tf_enable_onednn_opts=0

目前需要每次使用前单独执行这些命令；

——————

也可写入当前用户的~/.bashrc 文件的最后面

不过这样写还是没有通用性，一旦改为别的版本的 cuda 和 cudnn 以及 tensorrt的话，需要重新改写这些环境变量的值，特别是设置 tensorrt 的那句；

通用的方式可以使用软连接,创建一个类似/usr/local/cuda的软连接，在/usr/local/tensorrt,每次修改cuda的指向时，顺便修改tensorrt的指向；

4. 拷贝文件

在tensorrt的lib文件夹中，将 libnvinfer_plugin.so.7 拷贝成  libnvinfer_plugin.so.8

直接执行cp即可；因为tensorflow个别程序写的比较死，需要加载文件名为xxx.so.7的lib文件；

cp libnvinfer_plugin.so.8  libnvinfer_plugin.so.7

5. 安装python wheel

先查看python3的版本号：

python3 --version

返回3.8.10

进入wheel文件夹：

$ cd /usr/local/cuda-11.8/tensorrt_sdk/tensorrt-8.5.3.1/python
$ ls


 tensorrt-8.5.3.1-cp310-none-linux_x86_64.whl
 tensorrt-8.5.3.1-cp36-none-linux_x86_64.whl
 tensorrt-8.5.3.1-cp37-none-linux_x86_64.whl
 tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl
 tensorrt-8.5.3.1-cp39-none-linux_x86_64.whl

根据 python3的版本，本机只需要安装 tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl

版本对应：cp38  =》 3.8.10

$ sudo python3 -m pip install tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl

6. 安装uffwheel，为了 tensorflow中使用 tensorrt，

$ sudo python3 -m pip install uff-0.6.9-py2.py3-none-any.whl

类似这个样子：

7. 跑sample进行验证安装成功

因为要执行make，需要依赖 tensorrt的include/ *.h 文件; 官方采用的是相对路径；

这里有两个选择，

    1.  要么在用户文件夹中重新解压一次tensorrtxxx.tar.gz, 例如， /home/hanmeimei/ex_tensorrt/，然后跑到sample文件夹里 make

    2.  要么拷贝出来sample文件夹，修改里面的makefile.conf文件的变量，找到inlude文件夹，然后 make ；

执行效果：

tensorrt_sdk/tensorrt-8.5.3.1/bin$ ./sample_onnx_mnist

为了节省时间，这里直接在/usr/local/cuda-11.8/tensorrt_sdk/中 sudo make 并执行；

效果如图：

tensorrt 官方安装参考：

https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-tar

 而且 tensorrt是nvidia的一个开源加速库，repo在：

https://github.com/nvidia/tensorrt

——————————————————————————————————

七，重新测试tensorflow

保存并运行：line.py

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
 
trainx=np.linspace(-1,1,100)
trainy=2*trainx+np.random.randn(*trainx.shape)*0.3
#y=2x with noise
plt.plot(trainx,trainy,'ro',label='original data')
plt.legend()
plt.show()
 
 
x=tf.placeholder("float")
y=tf.placeholder("float")
w=tf.variable(tf.random_normal([1]),name="weight")
b=tf.variable(tf.zeros([1]),name="bias")
z=tf.multiply(x,w)+b
 
cost=tf.reduce_mean(tf.square(y-z))
learning_rate=0.01
optimizer=tf.train.gradientdescentoptimizer(learning_rate).minimize(cost)
 
trainepoch=20
displaystep=2
 
plotdata = { "batchsize": [] , "loss": [] }
def moving_average(a, w=10):
    if len(a) < w: 
        return a[:]    
    return [val if idx < w else sum(a[(idx-w):idx])/w for idx, val in enumerate(a)]
 
saver=tf.train.saver()
with tf.session() as sess:
    init=tf.global_variables_initializer()
    sess.run(init)
    
    for epoch in range(trainepoch):
        for (x,y) in zip(trainx,trainy):
            sess.run(optimizer,feed_dict={x:x,y:y})
 
        if epoch%displaystep==0:
            loss=sess.run(cost,feed_dict={x:trainx,y:trainy})
            print("epoch:",epoch+1,"cost=",loss,"w=",sess.run(w),"b=",sess.run(b))
            if not(loss=="na"):
                plotdata["batchsize"].append(epoch)
                plotdata["loss"].append(loss)
                
    print("finished!")
    print("cost=",sess.run(cost,feed_dict={x:trainx,y:trainy}),"w=",sess.run(w),"b=",sess.run(b))
    saver.save(sess,"d:/cuda/tensorflow_exam/study/3tensorflow-practice/the3rdmodel")

如果系统却类似matplotlib的话，使用pip3 安装：

$ sudo pip3 install matplotlib

运行：

$ python3 ./line.py