使用低级 TensorFlow API 进行分布式训练

数据挖掘 张量流 分散式
2022-02-16 12:10:45

我正在使用低级 Tensorflow API 进行模型训练。当我说低级别时,这意味着我正在定义tf.Session()图形的对象并在此会话中评估图形。

我想使用tf.distribute.MirroredStrategy().

我可以使用mirroredstrategy()tensorflow 在其文档中共享的示例在 tensorflow 顺序 API 上使用。

但是我在使用镜像策略执行 tf 低级代码时遇到了困难。

我尝试使用distribute.MirrorStrategy(),以下是资源利用的结果:

[0] GeForce RTX 2080 钛 | 48'C, 40 % | 10771 / 11019 MB | vipin(10763M) gdm(4M)
[1] GeForce RTX 2080 Ti | 37'C, 0 % | 10376 / 11014 MB | vipin(10327M) gdm(36M) gdm(8M)

即使模型使用了两个 GPU 的内存,但 GPU1 的利用率仍然为 0。

我不确定这个问题。甚至不确定 tensorflow 是否支持这一点。

请清除我的疑虑,并尽可能分享示例代码。

1个回答

1- 在代码开头(在库声明之后)使用distribute.MirrorStrategy() 以正确激活所有相关功能。

2-设置环境变量CUDA_VISIBLE_DEVICES=1(https://stackoverflow.com/questions/37893755/tensorflow-set-cuda-visible-devices-within-jupyter

3- 确保正确设置镜像变量,这要归功于:

with strategy.scope():
     x = tf.Variable(..)