TensorFlow 多卡训练 tf多卡训练

目录

export TF_GPU_ALLOCATOR=cuda_malloc_async

🔧 具体作用

优势

🧩 依赖条件

✅ 设置方式（Linux/macOS）

显存按需增长

---

export TF_GPU_ALLOCATOR=cuda_malloc_async

是设置 TensorFlow 使用 CUDA 异步内存分配器 的环境变量。这个设置可以带来性能上的优化，尤其在使用 GPU 的时候。

---

🔧 具体作用

TF_GPU_ALLOCATOR=cuda_malloc_async 的作用是：

让 TensorFlow 使用 CUDA 11.2+ 引入的 cudaMallocAsync API 来进行 GPU 显存的分配和释放。

---

优势

1. 提高性能：

   • 异步内存分配减少了 GPU 上的锁竞争。

   • 在多线程/并发数据加载和训练的情况下，效率明显提升。

2. 减少内存碎片：

   • cudaMallocAsync 使用内存池机制，比默认的 cudaMalloc 更智能地管理内存，减少碎片。

3. 更高效的资源利用：

   • 在一些模型（比如 Transformer）或动态图应用中，更容易避免 OOM（Out of Memory）。

---

🧩 依赖条件

• CUDA ≥ 11.2

• NVIDIA 驱动 ≥ 460

• TensorFlow ≥ 2.10（强烈建议 ≥ 2.11）

---

✅ 设置方式（Linux/macOS）

在终端或脚本中加上：

export TF_GPU_ALLOCATOR=cuda_malloc_async

然后运行你的训练脚本：

python train.py

调用代码：

export CUDA_VISIBLE_DEVICES=5,7

export TF_GPU_ALLOCATOR=cuda_malloc_async

python nlf/tf/main.py --train 

            with tf.device('/GPU:0'):
                loaded_model = tf.saved_model.load(FLAGS.load_backbone_from)

显存按需增长

for gpu in tf.config.experimental.list_physical_devices('GPU'):
    tf.config.experimental.set_memory_growth(gpu, True)
list_physical_devices('GPU')
列出当前可用的物理 GPU 设备。

set_memory_growth(gpu, True)
设置这个 GPU 为“按需增长显存”模式。也就是说，TensorFlow 只会根据需要逐步分配显存，而不是一启动就吃满整个 GPU。