需要事先安装较新版本的 cuda 和cudnn ，例如11.8 + 8.8
在已经安装过cuda的机器上安装新版cuda sdk 和 cudnn 可参考前述：
ubuntu 安装 多版本 cuda 11.4 11.8_Eloudy的博客-CSDN博客
一，安装python3 和 pip3
sudo apt-get install python3sudo apt-get install python3-pippip3 --versionsudo 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 jaxlibpip install --upgrade pip然后直接显式指定cuda和cudnn版本安装：
pip install "jax[cuda11_cudnn86]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html安装成功：
四，示例测试
python3
>>> import jax.numpy as np
>>>
import jax.numpy as npfrom jax import randomimport timex = 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：
生成三维数组
import jax.numpy as jnpfrom jax import grad, jit, vmapfrom jax import randomkey = random.PRNGKey(0)x = random.normal(key,(7, 2, 3))print(x)
计算 LU分解
import jax.numpy as jnpfrom jax import grad, jit, vmapfrom jax import randomimport pprintimport scipyimport scipy.linalgkey = 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)执行：
官方安装参考网址
五，顺便安装个tensorflow
先炸一炸最新版本
$ pip3 install tensorflow-gpu==33.33输出如下：
mklguy@mklguy--PC:~/ex/tensorflow_ex$ pip install tensorflow-gpu==33.33Defaulting to user installation because normal site-packages is not writeableERROR: 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
步骤：
下载 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
解压到对应cuda sdk版本的合适的位置
例如 /usr/local/cuda-11.8/
设置环境变量
export PATH=/usr/local/cuda/bin:$PATHexport LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATHexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.8/TensorRT_sdk/TensorRT-8.5.3.1/libexport TF_ENABLE_ONEDNN_OPTS=0目前需要每次使用前单独执行这些命令；
——————
也可写入当前用户的~/.bashrc 文件的最后面
不过这样写还是没有通用性，一旦改为别的版本的 cuda 和 cudnn 以及 TensorRT的话，需要重新改写这些环境变量的值，特别是设置 TensorRT 的那句；
通用的方式可以使用软连接,创建一个类似/usr/local/cuda的软连接，在/usr/local/tensorRT,每次修改cuda的指向时，顺便修改tensorRT的指向；
拷贝文件
在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安装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安装UFFwheel，为了 tensorflow中使用 TensorRT，
$ sudo python3 -m pip install uff-0.6.9-py2.py3-none-any.whl类似这个样子：
跑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
而且 TensorRT是Nvidia的一个开源加速库，repo在：
https://github.com/NVIDIA/TensorRT
——————————————————————————————————
七，重新测试tensorflow
保存并运行：line.py
import tensorflow as tfimport numpy as npimport matplotlib.pyplot as plt trainx=np.linspace(-1,1,100)trainy=2*trainx+np.random.randn(*trainx.shape)*0.3#y=2x with noiseplt.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.01optimizer=tf.train.GradientDescentOptimizer(learning_rate).minimize(cost) trainepoch=20displaystep=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
