tensorflow gpu 显存不足错误

Python

tensorflow默认会申请分配所有的gpu资源，这会导致明明可以跑的模型却无法跑起来，报OOM而无法创建sessions。

TensorFlow 在分配显存的时候，设置为可以逐步分配，需要多少显存，就一点点往上添，直到够用为止。比如GPU显存为8GB，开始训练一个模型，TensorFlow 先分配100MB，然后200MB 、300MB……最后在700MB处停住了，因为实际700MB就够用了。

  1. #指定使用那块GUP训练
  2. os.environ[“CUDA_VISIBLE_DEVICES”] = ‘0’
  3. config = tf.ConfigProto()
  4. # 设置最大占有GPU不超过显存的70%
  5. config.gpu_options.per_process_gpu_memory_fraction = 0.7
  6. # 重点：设置动态分配GPU
  7. config.gpu_options.allow_growth = True
  8. # 创建session时
  9. with tf.Session(config=config) as sess:

估算模型所需显存

一个模型占用显存主要包括：模型自身的参数、优化器参数、模型每层的输入输出。模型自身的参数主要指的是网络层的Weight和Bias。优化器参数指的是模型在反向传播过程中产生的参数，这部分参数主要是梯度。在SGD中，其大小与参数一样，因此模型的参数所占用的显存会翻倍。在Adam中，模型的参数量会翻4倍。

显存占用 = 模型显存占用 + batch_size × 每个样本的显存占用

通常如果模型的的参数按照float32（4个Byte）保存，那么显存的占用=参数量*4，但是实际需要的显存因为还有其它开销，所以比上述这种方式计算得到的值还要再大一些。

 

如果以上方法还存在问题？那就说明GPU显存是真的不够用了

  • 减少模型参数
  • 调低batchsize

本文链接地址: tensorflow gpu 显存不足错误

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注