tensorflow的显存分配机制

最新推荐文章于 2026-07-29 23:22:18 发布

原创

最新推荐文章于 2026-07-29 23:22:18 发布 · 713 阅读

收录于

当前文章被以下社区和专栏收录：

tensorflow的显存分配机制

默认情况下，TensorFlow 会映射进程可见的所有 GPU（取决于 CUDA_VISIBLE_DEVICES）的几乎全部内存。这是为了减少内存碎片，更有效地利用设备上相对宝贵的 GPU 内存资源。为了将 TensorFlow 限制为使用一组特定的 GPU，可以使用 tf.config.set_visible_devices 方法。

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only use the first GPU
  try:
    tf.config.set_visible_devices(gpus[0], 'GPU')
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus