Acme 在 Ubuntu22.04 上的执行 （gpu运行版本）

Acme 在 Ubuntu22.04 上的执行 （gpu运行版本）

0. WSL2 + Ubuntu22.04 的安装

如果没有 linux 双系统的话，可以安装 wsl2 + Ubuntu22.04 来作为替代方案。 网上教程很多，但其实直接 microsoft store 搜索安装就好了，安装结束后输入

wsl -l -v

检查是否安装成功，安装成功将出现如下的标识：

1. 按照 github 中 acme官网仓库的 readme中的installation 开始进行配置

官网仓库地址为：https://github.com/deepmind/acme
安装命令如下：

#创建虚拟环境
python3 -m venv acme
source acme/bin/activate
pip install --upgrade pip setuptools wheel
# 装jax 和 tensorflow
pip install dm-acme[jax,tf]
# 装envs
pip install dm-acme[envs]

如果能接受 cpu 的运行版本，差不多就结束了，如果要使用 gpu 应该还有一系列问题

2. 检查此虚拟环境下是否支持 gpu

我首先是在刚创建好的虚拟环境下运行了

# 进入python环境
python
# 尝试导入tensorflow
import tensorflow as tf

出现了如下报错

直接按照报错提示输入下面两行命令：

pip uninstall protobuf
pip install protobuf==3.20.1

再次 import ，发现没有报错了，但是仍然找不到 gpu ，内容如下所示：

于是在命令行中输入 nvidia-smi，发现自己根本没安装 nvidia，瞬间变成小丑，直接根据提示用apt 去 install，然后继续尝试 import tensorflow ，发现仍然有问题，报错和刚刚一样

tensorflow gpu 安装(ubuntu22.04)这篇博客完美解决了我的问题，值得一提的是，与博客不同的是 ， 你不需要再次执行类似 pip install tensorflow-gpu 等命令，因为在第1步配置环境的时候你已经安装了 tensorflow ，尽管你在pip list 中并没有看到你有所安装 tensorflow-gpu，但事实上，tf 在 后续版本中已经将 tensorflow 和 tensorflow-gpu结合了，对于 2.8.0版本，你是不需要额外进行安装的，这点非常重要！！！因为随意修改 tf 的版本可能会使你的 acme 无法运行

上面博客在运行到 add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /" 可能会出现报错：add-apt-repository: not found 错误，执行下面两条命令可以解决你的问题：

sudo apt-get update
sudo apt-get install software-properties-common

然后安装 cuda 结束后别忘记添加环境变量，如下所示：

# edit /etc/profile
sudo vim /etc/profile
# append path to "LD_LIBRARY_PATH" in profile file
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.8/targets/x86_64-linux/lib
# make environment file work
source /etc/profile

参考链接：StackOverflow

在执行完上述的所有内容之后，我 import tensorflow ， import 成功，且可以顺利地调用 gpu

3. jax 下能不能使用 gpu

在运行代码时，大概率会有这个问题：No GPU/TPU found, falling back to CPU
此时发现 tensorflow 可以检测到 gpu ，但是 jax 检测不到，检查了一圈，应该是没有配置 jaxlib

下面这个命令可以使你安装对应地 jax 和 jaxlib

pip install --upgrade "jax[cuda12_pip]==0.4.11" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

强烈建议安装版本为 0.4.11 ，因为官网地jax版本就是 0.4.11 ，如果不是 0.4.11 而是 0.4.14 的话，可能会出现如下报错：jax has no attribute xla 本人地解决方案比较逆天，谨慎参考，就是 cd 到 traceback 最后地那个报错地 python 代码文件夹，然后 ctrl+F ， 把 jax.xla 替换为 jax.interpreters.xla ，则问题可以被解决（记得只有一处被替换了，但其实这个替换是很有必要的）

在对 jax 动手的时候可能还会有很多其他问题，比如AttributeError: module 'jax.interpreters.xla' has no attribute 'DeviceArray' ， 这时候可以考虑更改 flax 的版本为 flax==0.6.11 且不更新它，可能会起到作用，参考的是这个链接的回答

这个过程可能还有很多细节我已经遗忘了，但大抵的思路是这样的，最后 jax 能不能用上 gpu 不用太强求，个人感觉速度上并没有太大的差距，但是为 tensorflow 配置 gpu 还是很有必要的