Keras限制GPU显存使用

深度学习门槛越来越低，尤其是Keras这样的高层次API加入以后，简单几行代码就能构建网络并得到不错的效果。

最近工作需要，开始使用Keras写3DConv做3D数据分类。

数据是肺部CT，目标是对检测网络得到的结果使用3D Conv网络做2分类。2D的检测网络获得肺部结节的ROI，但是由于二维上结节目标与肺部正常组织结构（如血管，气管等）特征类似，所以检测结果中包含了大量“假阳”，所以在检测网络的结果基础上使用3D网络，实际测试对“假阳”有很好抑制效果。

由于检测网络使用caffe框架，而3D网络又用以TensorFlow作为后端的Keras编写，结果TensorFlow跟caffe不兼容，同一个Python进程中不能同时包含这两个框架的环境。

于是将两部分代码分别放到不同的显卡上。TensorFlow如果不加限制会占满服务器上的所有显存，需要加以限制，查看了GitHub上的源码和解决方案如下：

首先是对keras配置文件进行修改，如果是在Linux环境下，也就是修改：

~/.keras/keras.json

这个文件。

修改内容如下：

{    "epsilon": 1e-07,    "floatx": "float32",    "image_data_format": "channels_last",    "backend": "tensorflow",    "gpu_options": {        "allow_growth": false,        "per_process_gpu_memory_fraction": 0.5,        "visible_device_list": "1"    }}

其中

per_process_gpu_memory_fraction

表示使用显存的比例；

visible_device_list

表示使用第几块显卡，0开始；

然后就需要修改后端代码：

在python库路径中找到keras位置，Linux下一般也就是：/usr/local/lib/python2.7/dist-packages/keras-2.****

这样一个路径，修改TensorFlow后端代码：

def get_session():    """Returns the TF session to be used by the backend.    If a default TensorFlow session is available, we will return it.    Else, we will return the global Keras session.    If no global Keras session exists at this point:    we will create a new global session.    Note that you can manually set the global session    via `K.set_session(sess)`.    # Returns        A TensorFlow session.    """    global _SESSION    if tf.get_default_session() is not None:        session = tf.get_default_session()    else:        if _SESSION is None:            _keras_base_dir = os.path.expanduser('~')            if not os.access(_keras_base_dir, os.W_OK):                 _keras_base_dir = '/tmp'            _keras_dir = os.path.join(_keras_base_dir, '.keras')            _config_path = os.path.expanduser(os.path.join(_keras_dir,                                                            'keras.json'))            if os.path.exists(_config_path):                try:                    _config = json.load(open(_config_path))                except ValueError:                    _config = {}            _options = _config.get('gpu_options', None)            _allow_growth = _options.get('allow_growth', False)            _mem_frac = _options.get('per_process_gpu_memory_fraction', 1.0)            _visible_device_list = _options.get('visible_device_list', None)            _gpu_options = tf.GPUOptions(allow_growth=_allow_growth,                                         per_process_gpu_memory_fraction=_mem_frac,                                         visible_device_list=_visible_device_list)            if not os.environ.get('OMP_NUM_THREADS'):                config = tf.ConfigProto(allow_soft_placement=True,                                        gpu_options=_gpu_options)            else:                num_thread = int(os.environ.get('OMP_NUM_THREADS'))                config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,                                        allow_soft_placement=True,                                        gpu_options=_gpu_options)            _SESSION = tf.Session(config=config)        session = _SESSION    if not _MANUAL_VAR_INIT:        with session.graph.as_default():            _initialize_variables()    return session

需要额外import 一个json库用来解析keras.json配置文件。

其实也就是对TensorFlow的session属性进行一下初始化。