1. 从一次显存溢出说起为什么GPU设置不是“装上驱动就完事”很多人第一次在深度学习项目里用上GPU心态都是这样的驱动装好了nvidia-smi能看到卡tf.test.is_gpu_available()返回True然后就觉得万事大吉直接开跑训练脚本。结果跑着跑着终端突然甩出一行红字ResourceExhaustedError: OOM when allocating tensor with shape [...]或者更隐蔽的情况——GPU利用率长期在10%以下晃悠训练速度跟CPU差不多白白浪费了硬件。这两种情况我都遇到过而且不止一次。问题往往不出在“GPU能不能用”而出在“GPU怎么用”。TensorFlow 对 GPU 的管理有一套自己的逻辑从设备可见性、显存分配策略到多卡并行、混合精度、算子放置每一层都有可以调的地方。默认配置能让你跑起来但未必能让你跑得好、跑得稳。这篇文章想聊的就是这件事在使用 TensorFlow 的过程中GPU 相关的设置到底有哪些关键点每个设置背后的逻辑是什么什么场景下该用哪种策略。不管你是刚接触GPU训练的新手还是已经跑过几个项目但总被显存问题困扰的开发者下面这些内容应该都能帮你少走一些弯路。我会尽量把每个参数、每个策略的“为什么”讲清楚而不是只丢一段代码让你复制。需要提前说明的是TensorFlow 1.x 和 2.x 在 GPU 管理上的 API 差异比较大本文以 TensorFlow 2.x 为主线涉及 1.x 的地方会单独标注。另外文中提到的所有操作都基于常见的单机多卡或单机单卡场景分布式多机训练不在讨论范围内。2. 先搞清楚TensorFlow眼里的GPU长什么样2.1 物理GPU、逻辑GPU与设备字符串的对应关系TensorFlow 在启动时会把系统里所有可用的 GPU 枚举一遍给它们分配逻辑设备编号。你通过tf.config.list_physical_devices(GPU)拿到的是物理设备列表返回的是一组PhysicalDevice对象每个对象有个name属性形如/physical_device:GPU:0、/physical_device:GPU:1。而实际在计算图中使用的设备字符串是/device:GPU:0这种形式注意这里没有physical前缀。这两个概念容易混淆。物理设备是硬件层面的逻辑设备是 TensorFlow 运行时层面的。默认情况下它们一一对应但你可以通过配置让一个物理 GPU 对应多个逻辑 GPU也可以让某些物理 GPU 对 TensorFlow 完全不可见。理解这一层区分是后面所有配置操作的基础。还有一个细节nvidia-smi里显示的 GPU 编号和 TensorFlow 里的GPU:0、GPU:1不一定是一一对应的。nvidia-smi的编号是驱动层面的TensorFlow 的编号是它自己枚举的顺序。在大多数单机场景下两者一致但如果你设置了CUDA_VISIBLE_DEVICES环境变量TensorFlow 看到的编号就会重新映射。比如你设置CUDA_VISIBLE_DEVICES2,3那么 TensorFlow 里的GPU:0实际上对应的是nvidia-smi里的 2 号卡GPU:1对应 3 号卡。这个映射关系在排查问题时非常关键很多人调了半天代码发现卡选错了就是没注意这一点。2.2 环境变量与运行时API两条控制路径的优先级控制 TensorFlow 使用哪些 GPU有两条路径环境变量和运行时 API。环境变量主要是CUDA_VISIBLE_DEVICES它在进程启动前就生效TensorFlow 初始化 CUDA 上下文时会读取这个变量只枚举其中指定的 GPU。运行时 API 则是tf.config.set_visible_devices()它在 TensorFlow 已经初始化之后调用可以动态调整可见设备列表。这两条路径有个重要的优先级关系CUDA_VISIBLE_DEVICES先起作用set_visible_devices是在前者的基础上再做筛选。举个例子假设机器上有 4 张卡你设置CUDA_VISIBLE_DEVICES0,1那么 TensorFlow 只能看到物理上的 0 和 1 号卡此时再调用tf.config.set_visible_devices([], GPU)就是把这两张卡也隐藏掉最终 TensorFlow 看不到任何 GPU。实际使用中如果是在脚本内部做设备管理用运行时 API 更灵活如果是在容器或集群调度环境里通常用环境变量更直接因为调度系统往往通过环境变量来分配资源。两种方式不要混用得太复杂否则排查起来很痛苦。注意tf.config.set_visible_devices()必须在任何 GPU 操作之前调用一旦 TensorFlow 已经初始化了 GPU 上下文比如你已经跑过一次模型训练再调用这个函数会报错。所以设备配置代码要放在脚本最前面import 完 TensorFlow 之后立刻执行。2.3 一个容易忽略的事实TensorFlow默认占满所有可见GPU这是新手最容易踩的坑之一。TensorFlow 2.x 默认的行为是进程启动后只要检测到可见的 GPU就会在第一次执行 GPU 算子时把每张卡的全部显存都预分配掉。注意是“全部显存”不是“按需分配”。这意味着如果你有一张 24GB 的卡跑一个只需要 2GB 显存的小模型TensorFlow 也会把 24GB 全部占住导致其他进程无法使用这张卡。这个设计有它的道理预分配可以避免训练过程中频繁的显存分配和释放带来的碎片化问题对性能有一定好处。但在实际开发中尤其是多人共用一台服务器、或者需要同时跑多个实验的场景下这个默认行为非常不友好。所以几乎所有的 TensorFlow GPU 使用指南都会告诉你第一件事就是设置显存分配策略。3. 显存分配策略从“独占整张卡”到“按需增长”3.1 为什么默认预分配全部显存会带来麻烦先具体说说预分配全部显存的问题。假设你在一台 8 卡服务器上做实验每张卡 24GB。你跑第一个实验TensorFlow 把 0 号卡的 24GB 全占了。这时候你想再跑一个对比实验指定用 1 号卡没问题。但如果你想在 0 号卡上再跑一个小任务就会直接 OOM即使第一个实验实际只用了 3GB 显存。更常见的情况是你跑完一个实验进程退出了但显存没有立刻释放CUDA 上下文销毁有延迟下一个实验启动时发现显存不够只能等或者手动清理。还有一种情况是 Jupyter Notebook 环境。你在 Notebook 里跑了一段训练代码然后想修改一下重新跑。如果之前的 TensorFlow 会话没有完全清理显存还占着新的会话就起不来。这种问题在开发阶段特别烦人因为你需要反复试错。所以设置显存分配策略不是“优化项”而是“必选项”。TensorFlow 提供了两种主要的显存分配模式按需增长和显存上限。3.2 按需增长模式的配置与适用场景按需增长memory growth的意思是TensorFlow 不再一次性预分配全部显存而是随着实际需求逐步向 GPU 申请。初始时只占很少一点随着模型和数据加载显存占用逐渐上升直到达到实际需要的量。配置方式如下import tensorflow as tf gpus tf.config.list_physical_devices(GPU) if gpus: try: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) except RuntimeError as e: print(e)这段代码要放在所有 GPU 操作之前。set_memory_growth是逐个 GPU 设置的不能一次性设置所有卡。设置之后你可以用nvidia-smi观察会发现 TensorFlow 进程的显存占用是逐步上升的而不是一上来就占满。按需增长适合什么场景开发调试阶段、多人共用服务器、需要同时跑多个小实验的情况。它的缺点是由于显存是动态申请的在训练过程中可能会有轻微的分配开销理论上对性能有一点点影响但在实际项目中这个影响通常可以忽略。另外按需增长模式下显存碎片化的问题会比预分配模式更明显一些如果模型需要频繁创建和销毁大量不同大小的张量可能会遇到“总显存够但连续显存不够”的情况。3.3 显存上限模式给每个进程划一条红线另一种策略是设置显存上限memory limit。你可以指定每个 GPU 上 TensorFlow 最多使用多少显存比如 8GB。这样即使卡有 24GBTensorFlow 也只会用到 8GB剩下的留给其他进程。gpus tf.config.list_physical_devices(GPU) if gpus: try: tf.config.set_logical_device_configuration( gpus[0], [tf.config.LogicalDeviceConfiguration(memory_limit8192)] ) except RuntimeError as e: print(e)注意这里的memory_limit单位是 MB8192 就是 8GB。这个配置也是必须在 GPU 初始化之前设置。显存上限模式适合什么场景当你明确知道模型需要多少显存并且希望严格控制资源使用时。比如在容器化部署中给每个容器分配固定的显存配额。或者在多人共用的服务器上给每个用户划定显存上限防止一个人把卡占满。但这里有个坑memory_limit设置的是逻辑设备的显存上限不是硬性隔离。如果设置了 8GB但模型实际需要 10GBTensorFlow 会在达到 8GB 时抛出 OOM而不是自动扩展到更多。所以设置这个值需要你对模型显存需求有比较准确的估计。3.4 两种策略的对比与选择建议对比维度按需增长显存上限配置方式set_memory_growth(gpu, True)set_logical_device_configuration显存占用动态上升按实际需求不超过设定上限适用场景开发调试、多人共用、多实验并行容器部署、资源配额管理性能影响轻微分配开销可忽略无明显影响主要风险碎片化导致OOM上限设太小导致OOM是否可组合可与显存上限组合使用可与按需增长组合使用实际使用中我个人的习惯是开发阶段一律开按需增长简单省事部署阶段根据模型实际显存需求设置上限留 20% 左右的余量。两种策略也可以组合先设置显存上限再开启按需增长这样既有上限保护又不会一次性占满。提示如果你在 Jupyter Notebook 里修改了显存配置需要重启 Kernel 才能生效。因为 TensorFlow 的 GPU 上下文一旦初始化配置就无法更改。这是很多人改了代码发现没效果的原因。4. 多卡场景下的设备放置与并行策略4.1 手动指定设备tf.device上下文管理器当你有多张 GPU 时TensorFlow 默认会把算子放在GPU:0上除非你显式指定其他设备。手动指定设备用tf.device上下文管理器with tf.device(/device:GPU:1): model build_model() model.fit(x_train, y_train)这段代码会把模型构建和训练都放在 1 号 GPU 上。tf.device可以嵌套也可以和tf.config.set_visible_devices配合使用。但手动指定设备有个问题它只影响算子的放置不改变数据分布。如果你想把训练数据也分到不同 GPU 上需要配合tf.distribute.Strategy来做。另外tf.device的作用范围是上下文内的所有算子。如果你在with块里调用了其他函数那些函数里的算子也会被放到指定设备上除非函数内部又用了tf.device覆盖。这个特性在组织代码时要注意避免不小心把不该放的算子放到 GPU 上。4.2 MirroredStrategy单机多卡最常用的并行方式MirroredStrategy是 TensorFlow 2.x 里单机多卡训练最常用的策略。它的工作原理是在每张 GPU 上复制一份完整的模型每张卡处理一部分数据然后通过 All-Reduce 操作同步梯度。这种方式也叫数据并行。strategy tf.distribute.MirroredStrategy() print(Number of devices: {}.format(strategy.num_replicas_in_sync)) with strategy.scope(): model build_model() model.compile(optimizeradam, losssparse_categorical_crossentropy)MirroredStrategy默认使用所有可见的 GPU。如果你想指定部分 GPU可以在创建策略时传入设备列表strategy tf.distribute.MirroredStrategy(devices[/gpu:0, /gpu:1])这里有个细节MirroredStrategy的devices参数接受的是逻辑设备字符串不是物理设备。如果你之前用set_visible_devices隐藏了某些卡这里的编号是隐藏后的编号。MirroredStrategy的显存管理有个特点每张卡上都会复制一份完整的模型参数和优化器状态。所以如果你的模型在单卡上需要 10GB 显存用 4 卡 MirroredStrategy 时每张卡仍然需要 10GB 左右加上梯度同步的额外开销。它不会把模型切分到多张卡上而是复制。这一点和模型并行不同选择策略时要清楚。4.3 梯度累积显存不够时的替代方案有时候你的模型太大单卡放不下但又没有多卡可用或者多卡也不够。这时候可以考虑梯度累积gradient accumulation。它的思路是用较小的 batch size 跑多次前向和反向累积梯度然后再更新一次参数。这样等效于用了更大的 batch size但显存占用只取决于小 batch。TensorFlow 本身没有内置的梯度累积 API需要手动实现。基本思路是用tf.GradientTape记录多次梯度然后累加accum_steps 4 optimizer tf.keras.optimizers.Adam() for step, (x_batch, y_batch) in enumerate(dataset): with tf.GradientTape() as tape: predictions model(x_batch, trainingTrue) loss loss_fn(y_batch, predictions) loss loss / accum_steps gradients tape.gradient(loss, model.trainable_variables) if (step 1) % accum_steps 0: optimizer.apply_gradients(zip(gradients, model.trainable_variables)) optimizer.zero_grad()这段代码的关键点是loss 要除以累积步数这样累积后的梯度才和真实大 batch 的梯度在量级上一致。另外要注意 BatchNorm 层的行为在梯度累积模式下BatchNorm 的统计量是基于小 batch 计算的可能和真实大 batch 有差异。如果模型里有 BatchNorm可能需要调整动量参数或者改用其他归一化方式。梯度累积的代价是训练速度变慢因为同样的数据量需要更多的前向反向次数。但它确实能让大模型在小显存上跑起来是一种用时间换空间的策略。4.4 混合精度训练用FP16省显存、提速度混合精度训练是另一个省显存、提速度的手段。它的核心思想是在前向和反向计算中使用 FP16半精度浮点数但参数更新时用 FP32单精度避免精度损失。TensorFlow 2.x 里开启混合精度比较简单from tensorflow.keras import mixed_precision policy mixed_precision.Policy(mixed_float16) mixed_precision.set_global_policy(policy)设置之后模型的大部分层会自动使用 FP16 计算。但有几类操作仍然会用 FP32Softmax、LayerNorm、损失计算等对数值稳定性敏感的操作。Keras 会自动处理这些细节。混合精度能省多少显存理论上 FP16 张量占用的显存是 FP32 的一半所以激活值和梯度相关的显存能省大约一半。但模型参数本身如果保持 FP32这部分显存不变。实际测试中整体显存节省通常在 30% 到 40% 之间取决于模型结构。速度提升方面在支持 Tensor Core 的 GPU 上比如 V100、A100、RTX 系列FP16 矩阵乘法的吞吐量是 FP32 的数倍训练速度提升明显。但在不支持 Tensor Core 的老卡上FP16 可能反而更慢因为需要额外的类型转换。注意开启混合精度后如果遇到 loss 变成 NaN 或者训练不稳定可以尝试调整 loss scale。Keras 的混合精度默认使用动态 loss scale大多数情况下能自动处理但如果问题持续可以手动设置mixed_precision.set_global_policy之前配置 loss scale 策略。5. 那些年我踩过的GPU配置坑5.1 坑一CUDA_VISIBLE_DEVICES设置顺序导致选错卡前面提到过CUDA_VISIBLE_DEVICES会重新映射 GPU 编号。我遇到过一个典型场景服务器上有 8 张卡0-3 号是同一批4-7 号是另一批。我想用 4-7 号卡跑实验于是设置CUDA_VISIBLE_DEVICES4,5,6,7。然后在代码里用tf.device(/gpu:0)以为用的是 4 号卡。结果nvidia-smi一看实际用的是 4 号卡没错但 TensorFlow 里的GPU:0对应的是物理 4 号GPU:1对应物理 5 号以此类推。这个映射本身没问题问题出在如果代码里硬编码了GPU:1而环境变量设置不同实际用的卡就会变。比如另一个人跑同样的代码但他设置的是CUDA_VISIBLE_DEVICES0,1那GPU:1就是物理 1 号卡。所以硬编码设备编号的代码在不同环境下行为不一致。我的建议是尽量用tf.config.list_physical_devices(GPU)动态获取设备列表然后根据列表长度和索引来分配而不是硬编码GPU:0、GPU:1。如果必须硬编码至少在代码开头打印一下实际可见的设备列表方便排查。5.2 坑二Jupyter Notebook里显存不释放在 Jupyter Notebook 里跑 TensorFlow 训练经常遇到一个问题跑完一个实验想修改代码重新跑结果报 OOM。明明之前的模型已经删了del model也执行了gc.collect()也调了显存还是占着。这个问题的根源是TensorFlow 的 GPU 上下文在 Notebook 进程里是全局的一旦初始化就不会销毁除非重启 Kernel。del model只能释放模型对象占用的显存但 TensorFlow 运行时本身占用的显存比如 CUDA 上下文、cuDNN 句柄等不会释放。解决办法有几个一是每次实验前重启 Kernel这是最彻底的二是用tf.keras.backend.clear_session()清理会话但这只能清理计算图不能完全释放显存三是把训练代码写成独立的 Python 脚本用subprocess调用这样每次运行都是新进程显存自然释放。我现在的习惯是Notebook 里只做数据探索和小规模调试正式训练一律写成脚本跑。这样既避免了显存问题也方便用nohup或tmux挂后台。5.3 坑三多进程同时访问GPU导致上下文冲突如果你用 Python 的multiprocessing做数据加载或并行处理要注意每个子进程都会尝试初始化自己的 CUDA 上下文。如果多个子进程同时访问同一张 GPU可能会遇到上下文冲突或者显存翻倍的问题。TensorFlow 的tf.data管道本身支持多进程数据加载它内部处理了 GPU 上下文的问题一般不会出问题。但如果你自己用multiprocessing.Process创建子进程并且在子进程里跑 TensorFlow 模型就需要小心。一个常见的做法是在子进程里设置CUDA_VISIBLE_DEVICES为空字符串让子进程只用 CPU。这样数据预处理在 CPU 上做主进程负责 GPU 计算。配置方式是在创建子进程之前设置环境变量或者在子进程的入口函数里设置。import os os.environ[CUDA_VISIBLE_DEVICES] import tensorflow as tf注意这个设置必须在 import TensorFlow 之前执行否则不生效。5.4 坑四显卡被其他进程占用但nvidia-smi看不到有时候nvidia-smi显示 GPU 显存被占用但看不到具体的进程信息只显示No running processes found。这种情况通常是之前的进程异常退出但显存没有释放。Linux 下可以用fuser -v /dev/nvidia*查看哪些进程占用了 GPU 设备文件然后手动 kill 掉。还有一种情况是容器里的进程占用了 GPU但在宿主机上nvidia-smi看不到进程名。这是因为容器内的进程命名空间和宿主机不同。这种情况下需要进入容器内部查看或者用nvidia-smi的--query-compute-apps参数查看计算进程。如果确认没有进程占用但显存仍然不释放可以尝试nvidia-smi --gpu-reset重置 GPU但这会影响该卡上的所有进程慎用。最稳妥的方式还是找到占用进程并正常结束它。6. 性能调优让GPU真正跑满6.1 输入管道优化别让CPU拖了GPU的后腿GPU 利用率低很多时候不是 GPU 本身的问题而是数据供给跟不上。GPU 算得飞快但 CPU 还在慢慢读文件、做预处理GPU 大部分时间在等数据。这种情况在图像、视频类任务中特别常见。TensorFlow 的tf.dataAPI 提供了一套输入管道优化手段。核心思路是用prefetch让数据准备和模型计算重叠用map的num_parallel_calls参数并行做预处理用cache缓存已经处理过的数据避免重复计算。dataset tf.data.Dataset.from_tensor_slices((x_train, y_train)) dataset dataset.map(preprocess_fn, num_parallel_callstf.data.AUTOTUNE) dataset dataset.cache() dataset dataset.shuffle(buffer_size10000) dataset dataset.batch(batch_size) dataset dataset.prefetch(buffer_sizetf.data.AUTOTUNE)AUTOTUNE让 TensorFlow 根据可用 CPU 核心数自动决定并行度大多数情况下够用。如果数据量太大cache放不下可以改用cache到文件或者只用prefetch。判断输入管道是否是瓶颈可以看nvidia-smi里的 GPU 利用率。如果利用率长期低于 70%而且波动很大大概率是数据供给问题。另一个方法是看 TensorFlow 的 profiler 输出它会告诉你每个算子的耗时分布。6.2 XLA编译用JIT换性能XLAAccelerated Linear Algebra是 TensorFlow 的一个编译器后端它可以把多个算子融合成一个减少 kernel 启动开销和内存访问。开启 XLA 很简单tf.config.optimizer.set_jit(True)或者在model.compile时指定jit_compileTruemodel.compile(optimizeradam, loss..., jit_compileTrue)XLA 对某些模型结构提升明显尤其是包含大量小算子的模型。但它也不是万能的编译本身需要时间如果模型结构频繁变化比如动态 shapeXLA 可能反而拖慢速度。另外XLA 对某些自定义算子的支持不完善可能会报错。我的建议是先不开 XLA 跑一遍记录 baseline 性能然后开 XLA 再跑一遍对比速度和显存占用。如果提升明显且没有报错就保留如果没提升或者报错就关掉。不要盲目开 XLA它不是一个“开了就快”的开关。6.3 算子放置的隐式规则与显式控制TensorFlow 在放置算子时有一套隐式规则如果算子有 GPU 实现默认放在 GPU 上如果没有放在 CPU 上。但有些算子虽然有 GPU 实现放在 CPU 上可能更合适比如小规模的索引操作、形状变换等。你可以通过tf.config.set_soft_device_placement(True)开启软放置这样当 GPU 上没有对应算子实现时TensorFlow 会自动回退到 CPU而不是报错。这个选项在跨设备兼容性上很有用但可能会掩盖一些性能问题。显式控制算子放置用tf.device前面已经讲过。这里补充一点tf.device的作用是“建议”而非“强制”如果指定设备上没有算子实现TensorFlow 仍然会回退到 CPU在软放置开启的情况下。所以不要以为用了tf.device(/gpu:0)就一定能跑在 GPU 上关键还是看算子本身有没有 GPU 实现。7. 写在最后一些个人习惯上面聊了很多配置和策略最后说几个我在实际项目中养成的习惯供参考。第一任何 TensorFlow 脚本的开头我都会加一段设备信息打印import tensorflow as tf print(TensorFlow version:, tf.__version__) gpus tf.config.list_physical_devices(GPU) print(GPUs:, gpus) for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)这几行代码花不了多少时间但能在出问题时快速定位是环境问题还是代码问题。第二训练脚本一定用nohup或tmux挂后台输出重定向到日志文件。这样即使 SSH 断开训练也不会中断而且日志可以事后分析。第三显存配置不要写死在代码里用环境变量或配置文件控制。这样同一份代码可以在不同机器、不同卡数上运行不需要改代码。第四遇到 OOM 不要急着加卡或换大卡先分析显存都用在哪了。TensorFlow 的 profiler 可以给出显存占用的详细分布很多时候优化一下 batch size 或输入管道就能解决不需要动硬件。GPU 设置这件事说复杂也复杂说简单也简单。核心就是搞清楚 TensorFlow 的设备管理逻辑然后根据实际场景选择合适的策略。希望上面这些内容能帮你少踩几个坑。