应用场景:

      在实验室或者公司时,一般会遇到服务器紧张的情况。在这种情况下,可以通过限制GPU所占显存来使得一台服务器可以同时跑几个模型,当然,土豪实验室或者公司不用担心这个问题。

由于暂时只接触了Caffe,Tensorflow和keras,现只提供这三种深度学习框架的控制方法。

解决办法:

      对于caffe而言,默认不是占满GPU全部的显存,而是通过网络设置batchsize及每个epoch所迭代的次数,网络模型大小等参数控制,在运行时遇到超出内存等报错时,可以调低batchsize数量及网络模型的大小。

     caffe:调整网络参数及配置文件参数以及选择对应的GPU/CPU

batch_size: 64
iter_size: 2


# CPU
将examples中你想运行的例子对应的prototxt文件中的GPU更改为CPU

     对于tensorflow及keras而言,二者默认都是占用GPU全部的内存,由于keras是集成在tensorflow中,因此,二者的解决办法类似。

    在文件中添加一些命令以控制模型所占用GPU显存的比例。

     Tensorflow:添加几行语句即可

# init tf session
tf_config = tf.ConfigProto()
if self.config['GPU_MEN_FRACTION'] > 0.0:
    tf_config.gpu_options.per_process_gpu_memory_fraction = 0.3  # 这里设置占用的比例
else:
    # disable gpu memory allocation
    tf_config.device_count['GPU'] = 0
    # tf_config.device_count = {'GPU': 0}
self.sess = tf.Session(config=tf_config)

  同样,使用tensorflow实现keras也是添加一些命令即可控制所占的比例。

Keras: 添加几行语句即可,GPU/CPU

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
config.gpu_options.visible_device_list = "0"
set_session(tf.Session(config=config))
# only use cpu
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = ""

Pytorch: 添加几行语句即可,GPU/CPU

# 配置GPU,这里有两种方式
## 方案一:使用os.environ
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 方案二:使用“device”,后续对要使用GPU的变量用.to(device)即可
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

参考资料:

  1. Limit the resource usage for tensorflow backend

  2. Is there any ways to reduce the GPU Memory Caffe use?      

  3. https://www.tensorflow.org/guide/gpu?hl=zh-cn

更多推荐