src/ports/postgres/modules/deep_learning/gpu_info_from_tf.py_in

# coding=utf-8
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""
@file gpu_info_from_tf.py_in

@brief This file prints out the gpu descriptions using tensorflow API. This file
is intended to be called using subprocess. See madlib_keras_gpu_info.py_in
for more details.
"""

import tensorflow as tf
from tensorflow.python.client import device_lib
from tensorflow.keras import backend as K

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
local_device_protos = device_lib.list_local_devices()
K.clear_session()
sess.close()
if local_device_protos:
    for x in local_device_protos:
        if x.device_type == 'GPU':
            print x.physical_device_desc