文 / 李锡涵,Google Developers Expert
本文节选自《简单粗暴 TensorFlow 2.0》
![a792efc299bbf4c43764cb11acc49c5c.gif](https://i-blog.csdnimg.cn/blog_migrate/790bebf82ffb794268f5ab3ed02dcfad.gif)
在上一篇文章中,我们介绍了 tf.TensorArray 的建立方式。
本来接下来将介绍 TensorFlow 中的常用模块 tf.config,如何合理分配显卡资源,设置显存使用策略。
指定当前程序使用的 GPU
很多时候的场景是:实验室 / 公司研究组里有许多学生 / 研究员需要共同使用一台多 GPU 的工作站,而默认情况下 TensorFlow 会使用其所能够使用的所有 GPU,这时就需要合理分配显卡资源。
首先,通过tf.config.experimental.list_physical_devices
,我们可以获得当前主机上某种特定运算设备类型(如
GPU
或
CPU
)的列表,例如,在一台具有 4 块 GPU 和一个 CPU 的工作站上运行以下代码:
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)
输出:
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'),
PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU'),
PhysicalDevice(name='/physical_device:GPU:2', dev