TensorFlow 多卡训练 tf多卡训练

目录

[export TF_GPU_ALLOCATOR=cuda_malloc_async](#export TF_GPU_ALLOCATOR=cuda_malloc_async)

[🔧 具体作用](#🔧 具体作用)

优势

[🧩 依赖条件](#🧩 依赖条件)

[✅ 设置方式（Linux/macOS）](#✅ 设置方式（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)