gpu_utils.py

def parallel_gpu_jobs(allow_growth=True, fraction=.5):

    '''Sets the max used memory as a fraction for tensorflow
    backend

    allow_growth :: True of False

    fraction :: a float value (e.g. 0.5 means 4gb out of 8gb)

    '''

    import keras.backend as K
    import tensorflow as tf

    gpu_options = K.tf.GPUOptions(allow_growth=allow_growth,
                                  per_process_gpu_memory_fraction=fraction)
    config = tf.ConfigProto(gpu_options=gpu_options)
    session = K.tf.Session(config=config)
    K.set_session(session)


def multi_gpu(model, gpus=None, cpu_merge=True, cpu_relocation=False):

    '''Takes as input the model, and returns a model
    based on the number of GPUs available on the machine
    or alternatively the 'gpus' user input.

    NOTE: this needs to be used before model.compile() in the
    model inputted to Scan in the form:

    from talos.utils.gpu_utils import multi_gpu
    model = multi_gpu(model)

    '''

    from keras.utils import multi_gpu_model

    return multi_gpu_model(model,
                           gpus=gpus,
                           cpu_merge=cpu_merge,
                           cpu_relocation=cpu_relocation)


def force_cpu():

    '''Force CPU on a GPU system
    '''

    import keras.backend as K
    import tensorflow as tf

    config = tf.ConfigProto(device_count={'GPU': 0})
    session = tf.Session(config=config)
    K.set_session(session)