应用场景:
在实验室或者公司时,一般会遇到服务器紧张的情况。在这种情况下,可以通过限制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