深度学习|Keras限制tf后端的gpu显存用量

  • 训练模型
## keras example imports from keras.models import Sequential from keras.layers import Dense, Dropout from keras.layers import Embedding from keras.layers import LSTM ## extra imports to set GPU options import tensorflow as tf from keras import backend as k ################################### # TensorFlow wizardry config = tf.ConfigProto() # Don't pre-allocate memory; allocate as-needed config.gpu_options.allow_growth = True # Only allow a total of half the GPU memory to be allocated config.gpu_options.per_process_gpu_memory_fraction = 0.5 # Create a session with the above options specified. k.tensorflow_backend.set_session(tf.Session(config=config)) ################################### model = Sequential() model.add(Embedding(max_features, output_dim=256)) model.add(LSTM(128)) model.add(Dropout(0.5)) model.add(Dense(1, activation='sigmoid')) model.compile(loss='binary_crossentropy', optimizer='rmsprop', metrics=['accuracy']) model.fit(x_train, y_train, batch_size=16, epochs=10) score = model.evaluate(x_test, y_test, batch_size=16)

  • 加载模型
# keras example imports from keras.models import load_model ## extra imports to set GPU options import tensorflow as tf from keras import backend as k ################################### # TensorFlow wizardry config = tf.ConfigProto() # Don't pre-allocate memory; allocate as-needed config.gpu_options.allow_growth = True # Only allow a total of half the GPU memory to be allocated config.gpu_options.per_process_gpu_memory_fraction = 0.5 # Create a session with the above options specified. k.tensorflow_backend.set_session(tf.Session(config=config)) ################################### # returns a compiled model # identical to the previous one model = load_model('my_model.h5') # TODO: classify all the things

【深度学习|Keras限制tf后端的gpu显存用量】参考:
https://michaelblogscode.wordpress.com/2017/10/10/reducing-and-profiling-gpu-memory-usage-in-keras-with-tensorflow-backend/

    推荐阅读