search

PPPLDeepLearning / plasma-python

PPPL deep learning disruption prediction package
http://tigress-web.princeton.edu/~alexeys/docs-web/html/
79 stars 43 forks source link

Tensorflow backend #7

Closed ASvyatkovskiy closed 7 years ago

ASvyatkovskiy commented 7 years ago

The purpose of this PR is to enable TensorFlow backend support in FRNN. The version has been tested on a single GPU on Tigressdata, full multi-GPU support will follow in subsequent pull requests after migration to Keras-2.0 is finished.

Key changes: Control backend with a single parameter in the FRNN conf.yaml

model:
    backend: 'tensorflow'

This will modify the Keras backend at runtime:

if backend == 'tf' or backend == 'tensorflow':
    os.environ['KERAS_BACKEND'] = 'tensorflow'

Use TensorFlow ConfigProto: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/config.proto https://www.tensorflow.org/api_docs/python/tf/ConfigProto

    import tensorflow
    from keras.backend.tensorflow_backend import set_session
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.95, allow_growth=True)
    config = tf.ConfigProto(gpu_options=gpu_options)
    set_session(tf.Session(config=config))

Currently, GPU devices is controlled via CUDA_VISIBLE_DEVICES env variable. Consider a possibility of using visible_device_list