TensorFlow GPU memory?

Question

Python

I am training a neural network on the GPU (TensorFlow + Keras), the problem is that after running the script, the GPU memory overflows and I get an error:

E tensorflow/stream_executor/cuda/cuda_driver.cc:806] failed to allocate 4.61G (4946656768 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY: out of memory

Here is the script code:

from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf
from tensorflow import keras

import numpy as np
import matplotlib.pyplot as plt

fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

train_images = train_images / 255.0

test_images = test_images / 255.0

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(800, activation="relu"),
    keras.layers.Dense(10, activation="softmax")
])

model.compile(optimizer=tf.train.AdamOptimizer(),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=100, verbose=1, batch_size=200)

Answer 1

The solution is pretty simple.
Enough:

from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf
from tensorflow import keras

import numpy as np
import matplotlib.pyplot as plt

fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

train_images = train_images / 255.0

test_images = test_images / 255.0

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(800, activation="relu"),
    keras.layers.Dense(10, activation="softmax")
])
opt = tf.keras.optimizers.Adam(lr=0.001)
model.compile(optimizer=opt,
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=100, verbose=1, batch_size=200)

Add lines before working with TensorFlow:

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

Answer 2

Problem solved.