仅支持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。

06-19 19:03