GPU memory runing out on Training a neural network on MNIST with Keras

/!\ IF YOU WANT PEOPLE TO HELP YOU, PLEASE GIVE AS MUCH DETAIL AS POSSIBLE, INCLUDING THE FULL STACKTRACE AND CODE SNIPPET
2021-08-21 20:17:49.231141: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
2021-08-21 20:17:50.302539: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcuda.so.1
2021-08-21 20:17:50.322448: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-21 20:17:50.322648: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties:
pciBusID: 0000:01:00.0 name: NVIDIA GeForce GTX 1080 computeCapability: 6.1
coreClock: 1.7715GHz coreCount: 20 deviceMemorySize: 7.93GiB deviceMemoryBandwidth: 298.32GiB/s
2021-08-21 20:17:50.322662: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
2021-08-21 20:17:50.325336: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublas.so.11
2021-08-21 20:17:50.325363: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcublasLt.so.11
2021-08-21 20:17:50.326597: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcufft.so.10
2021-08-21 20:17:50.326809: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcurand.so.10
2021-08-21 20:17:50.327222: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcusolver.so.11
2021-08-21 20:17:50.327861: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcusparse.so.11
2021-08-21 20:17:50.327980: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudnn.so.8
2021-08-21 20:17:50.328059: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-21 20:17:50.328360: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-21 20:17:50.328533: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
2021-08-21 20:17:50.328794: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-08-21 20:17:50.329199: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-21 20:17:50.329380: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties:
pciBusID: 0000:01:00.0 name: NVIDIA GeForce GTX 1080 computeCapability: 6.1
coreClock: 1.7715GHz coreCount: 20 deviceMemorySize: 7.93GiB deviceMemoryBandwidth: 298.32GiB/s
2021-08-21 20:17:50.329434: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-21 20:17:50.329605: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2021-08-21 20:17:50.329769: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
2021-08-21 20:17:50.329804: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
Traceback (most recent call last):
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/data/util/structure.py", line 106, in normalize_element
spec = type_spec_from_value(t, use_fallback=False)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/data/util/structure.py", line 480, in type_spec_from_value
(element, type(element).name))
TypeError: Could not build a TypeSpec for ['/home/vincent/tensorflow_datasets/mnist/3.0.1/mnist-train.tfrecord-00000-of-00001'] with type list
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/usr/lib/python3.6/code.py", line 91, in runcode
exec(code, self.locals)
File "", line 1, in
File "/home/vincent/Applications/pycharm-2021.2/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 198, in runfile
pydev_imports.execfile(filename, global_vars, local_vars) # execute the script
File "/home/vincent/Applications/pycharm-2021.2/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "/home/vincent/PycharmProjects/federated_test/nonFederated_lenet5.py", line 9, in
with_info=True,
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/load.py", line 330, in load
ds = dbuilder.as_dataset(**as_dataset_kwargs)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/logging/init.py", line 81, in decorator
return function(*args, **kwargs)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/dataset_builder.py", line 563, in as_dataset
datasets = utils.map_nested(build_single_dataset, split, map_tuple=True)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/utils/py_utils.py", line 176, in map_nested
map_nested(function, v, dict_only, map_tuple) for v in data_struct
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/utils/py_utils.py", line 176, in
map_nested(function, v, dict_only, map_tuple) for v in data_struct
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/utils/py_utils.py", line 183, in map_nested
return function(data_struct)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/dataset_builder.py", line 585, in _build_single_dataset
read_config=read_config,
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/dataset_builder.py", line 938, in _as_dataset
disable_shuffling=self.info.disable_shuffling,
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/tfrecords_reader.py", line 417, in read
return tf.nest.map_structure(_read_instruction_to_ds, instructions)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 867, in map_structure
structure[0], [func(*x) for x in entries],
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 867, in
structure[0], [func(*x) for x in entries],
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/tfrecords_reader.py", line 414, in _read_instruction_to_ds
decode_fn=decode_fn,
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/tfrecords_reader.py", line 460, in read_files
file_format=self._file_format,
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow_datasets/core/tfrecords_reader.py", line 256, in read_files
instruction_ds = tf.data.Dataset.from_tensor_slices(tensor_inputs)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 760, in from_tensor_slices
return TensorSliceDataset(tensors)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 3320, in init
element = structure.normalize_element(element)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/data/util/structure.py", line 111, in normalize_element
ops.convert_to_tensor(t, name="component%d" % i))
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/profiler/trace.py", line 163, in wrapped
return func(*args, **kwargs)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1566, in convert_to_tensor
ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 339, in _constant_tensor_conversion_function
return constant(v, dtype=dtype, name=name)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 265, in constant
allow_broadcast=True)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 276, in _constant_impl
return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 301, in _constant_eager_impl
t = convert_to_eager_tensor(value, ctx, dtype)
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/framework/constant_op.py", line 97, in convert_to_eager_tensor
ctx.ensure_initialized()
File "/home/vincent/venv/lib/python3.6/site-packages/tensorflow/python/eager/context.py", line 525, in ensure_initialized
context_handle = pywrap_tfe.TFE_NewContext(opts)
tensorflow.python.framework.errors_impl.InternalError: CUDA runtime implicit initialization on GPU:0 failed. Status: out of memory

Short description
Description of the bug.

I used pycharm to run this demo. But I can not run it a second time as the GPU memory is fully occupied after the first running.

Environment information

• Operating System: ubuntu 18

• Python version: 3.6

• tensorflow-datasets/tfds-nightly version:

• tensorflow/tf-nightly version:

• Does the issue still exists with the last tfds-nightly package (pip install --upgrade tfds-nightly) ?

Reproduction instructions

<put a code snippet or link to a gist here>

Link to logs
If applicable, <link to gist with logs, stack trace>

Expected behavior
What you expected to happen.

Additional context
Add any other context about the problem here.

The failure is reported from /home/vincent/PycharmProjects/federated_test/nonFederated_lenet5.py while loading MNIST through tensorflow_datasets. Start by obtaining the missing code, package versions, and a reproducible run, then inspect the TensorFlow and TFDS loading path shown in the traceback. Done means determining whether the failure is a TFDS defect or an environment/configuration issue and documenting a verified reproduction or resolution.