官网方法，conda环境安装tensorflow，可使用多个GPU

最新推荐文章于 2024-08-02 21:11:23 发布

LRJ-jonas

最新推荐文章于 2024-08-02 21:11:23 发布

阅读量926

点赞数

文章标签： tensorflow python 人工智能深度学习 conda

本文链接：https://blog.csdn.net/m0_55097528/article/details/131288972

版权

文章介绍了如何使用conda创建支持CPU和GPU的TensorFlow环境。对于GPU，它强调了CUDA和cudnn的版本要求，并指出从TensorFlow2.0开始需要CUDA10.0以上。TensorFlow2.4引入的tf.distribute.MirroredStrategy用于单机多GPU的同步训练，使得模型变量在GPU间镜像同步。

摘要由CSDN通过智能技术生成

仅支持CPU的方案：

conda create -n tf tensorflow
conda activate tf

创建一个名为‘tf’的conda环境，自动安装cpu的tensorflow

可以使用gpu的方案：

conda create -n tf-gpu tensorflow-gpu
conda activate tf-gpu

使用conda安装 GPU TensorFlow 会安装依赖库：CUDA（cuda toolkit）以及cudnn。（按理说conda会自动安装好兼容的版本）

对于 Windows 和 Linux 系统，CUDA 10.0 以后的版本才支持 TensorFlow 2.0 以后的版本。

无需更改任何代码，TensorFlow 代码以及 tf.keras 模型就可以在单个 GPU 上透明运行。

注：使用 tf.config.list_physical_devices('GPU') 可以确认 TensorFlow 使用的是 GPU。

在一台或多台机器上，要顺利地在多个 GPU 上运行，最简单的方法是使用分布策略。

在 TensorFlow 2.4 版本中，可以使用 tf.distribute.MirroredStrategy 来在一台机器的多个 GPU 上进行同步分布式训练。该策略会为每个 GPU 设备创建一个副本。模型中的每个变量都会在所有副本之间进行镜像。这些变量将共同形成一个名为 MirroredVariable 的单个概念变量。这些变量会通过应用相同的更新来保持同步。1

以下是一个使用 tf.distribute.MirroredStrategy 的示例代码：

mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
    model = create_model()
model.fit(train_dataset, epochs=10, steps_per_epoch=100)

请注意，tf.distribute.MirroredStrategy 只能在一台机器上的多个 GPU 上进行同步分布式训练。如果您想要在多台机器上进行分布式训练，您可以使用其他分布式策略，例如 tf.distribute.experimental.MultiWorkerMirroredStrategy 或 tf.distribute.experimental.TPUStrategy。