TensorFlow 读取GPU并设置内存自增长

TensorFlow 读取GPU并设置内存自增长

import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sklearn
import pandas as pd
import os
import sys
import time
import tensorflow as tf

from tensorflow_core.python.keras.api._v2 import keras

print(tf.__version__)
print(sys.version_info)
for module in mpl, np, pd, sklearn, tf, keras:
    print(module.__name__, module.__version__)

# 打印模型变量所在位置（GPU，CPU）
tf.debugging.set_log_device_placement(True)  

# 获取物理GPU的个数
gpus = tf.config.experimental.list_physical_devices("GPU")  
for gpu in gpus:
    # 设置内存增长方式 自增长
    tf.config.experimental.set_memory_growth(gpu, True)  
print("物理GPU个数:", len(gpus))

# 获取逻辑GPU的个数
logical_gpus = tf.config.experimental.list_logical_devices("GPU") 
print("逻辑GPU个数:", len(logical_gpus))

tensorflow 1.13 GPU + keras 内存自增长

import numpy as np
import keras.backend as K
import tensorflow as tf

# 关键代码
# 设置自增长
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# 设置到keras 中去
K.tensorflow_backend.set_session(tf.Session(config=config))