Tensorflow doesn't work with gpu - too much memory is used. How to solve it?

Question

I use tensorflow for image classification (20 classes) with convolutions. My dataset contains about 20000 train images and 5000 test images. Images (RGB) have 200x256 pixels. When I run script to train model using cpu everything seems fine. However when I try run script using gpu, after loading my training and test data I get error on model_fit function.

Num GPUs Available:  1
2022-05-04 17:58:58.482057: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX AVX2
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-05-04 17:59:03.655618: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 4634 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1060, pci bus id: xxxx:xx:xx.x, compute capability: 6.1
1 Physical GPUs, 1 Logical GPUs
Path: D:/Dataset/seg_train
Loading seg_train
Path: D:/Dataset/seg_test
Loading seg_test
2022-05-04 18:02:48.971100: W tensorflow/core/common_runtime/bfc_allocator.cc:462] Allocator (GPU_0_bfc) ran out of memory trying to allocate 10.44GiB (rounded to 11206656000)requested by op _EagerConst
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation.
Current allocation summary follows.
Current allocation summary follows.
2022-05-04 18:02:48.996013: I tensorflow/core/common_runtime/bfc_allocator.cc:1010] BFCAllocator dump for GPU_0_bfc
2022-05-04 18:02:48.996173: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (256):  Total Chunks: 16, Chunks in use: 16. 4.0KiB allocated for chunks. 4.0KiB in use in bin. 392B client-requested in use in bin.
2022-05-04 18:02:48.996308: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (512):  Total Chunks: 1, Chunks in use: 1. 512B allocated for chunks. 512B in use in bin. 512B client-requested in use in bin.
2022-05-04 18:02:48.996473: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (1024):         Total Chunks: 1, Chunks in use: 1. 1.2KiB allocated for chunks. 1.2KiB in use in bin. 1.0KiB client-requested in use in bin.
2022-05-04 18:02:48.996629: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (2048):         Total Chunks: 2, Chunks in use: 1. 7.0KiB allocated for chunks. 3.5KiB in use in bin. 3.4KiB client-requested in use in bin.
2022-05-04 18:02:48.996889: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (4096):         Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:48.997493: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (8192):         Total Chunks: 1, Chunks in use: 1. 9.5KiB allocated for chunks. 9.5KiB in use in bin. 9.5KiB client-requested in use in bin.
2022-05-04 18:02:48.997960: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (16384):        Total Chunks: 1, Chunks in use: 0. 19.0KiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:48.998482: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (32768):        Total Chunks: 2, Chunks in use: 1. 79.5KiB allocated for chunks. 36.0KiB in use in bin. 36.0KiB client-requested in use in bin.
2022-05-04 18:02:48.999113: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (65536):        Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:48.999710: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (131072):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.000273: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (262144):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.000742: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (524288):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.001208: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (1048576):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.001671: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (2097152):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.002131: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (4194304):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.002700: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (8388608):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.004034: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (16777216):     Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.004682: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (33554432):     Total Chunks: 1, Chunks in use: 1. 44.56MiB allocated for chunks. 44.56MiB in use in bin. 44.56MiB client-requested in use in bin.
2022-05-04 18:02:49.005383: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (67108864):     Total Chunks: 1, Chunks in use: 0. 89.12MiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.007520: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (134217728):    Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.008016: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (268435456):    Total Chunks: 1, Chunks in use: 0. 4.39GiB allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2022-05-04 18:02:49.008477: I tensorflow/core/common_runtime/bfc_allocator.cc:1033] Bin for 10.44GiB was 256.00MiB, Chunk State:
2022-05-04 18:02:49.008888: I tensorflow/core/common_runtime/bfc_allocator.cc:1039]   Size: 4.39GiB | Requested Size: 0B | in_use: 0 | bin_num: 20, prev:   Size: 44.56MiB | Requested Size: 44.56MiB | in_use: 1 | bin_num: -1
2022-05-04 18:02:49.009335: I tensorflow/core/common_runtime/bfc_allocator.cc:1046] Next region of size 4859428864
2022-05-04 18:02:49.025604: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00000 of size 256 next 1
2022-05-04 18:02:49.025772: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00100 of size 1280 next 2
2022-05-04 18:02:49.026373: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00600 of size 256 next 3
2022-05-04 18:02:49.026991: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00700 of size 256 next 4
2022-05-04 18:02:49.028407: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00800 of size 256 next 5
2022-05-04 18:02:49.028560: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00900 of size 256 next 6
2022-05-04 18:02:49.029196: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00a00 of size 256 next 9
2022-05-04 18:02:49.029937: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00b00 of size 256 next 10
2022-05-04 18:02:49.030556: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00c00 of size 256 next 11
2022-05-04 18:02:49.031054: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00d00 of size 256 next 14
2022-05-04 18:02:49.031553: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00e00 of size 256 next 15
2022-05-04 18:02:49.031906: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a00f00 of size 512 next 16
2022-05-04 18:02:49.032334: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a01100 of size 256 next 19
2022-05-04 18:02:49.032719: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a01200 of size 256 next 20
2022-05-04 18:02:49.033158: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a01300 of size 256 next 21
2022-05-04 18:02:49.033523: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a01400 of size 256 next 24
2022-05-04 18:02:49.033892: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a01500 of size 256 next 25
2022-05-04 18:02:49.034323: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a01600 of size 256 next 26
2022-05-04 18:02:49.034824: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] Free  at b03a01700 of size 3584 next 7
2022-05-04 18:02:49.035472: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a02500 of size 3584 next 8
2022-05-04 18:02:49.035923: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] Free  at b03a03300 of size 19456 next 23
2022-05-04 18:02:49.036957: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a07f00 of size 9728 next 22
2022-05-04 18:02:49.039251: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] Free  at b03a0a500 of size 44544 next 13
2022-05-04 18:02:49.039789: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b03a15300 of size 36864 next 12
2022-05-04 18:02:49.040234: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] Free  at b03a1e300 of size 93454336 next 18
2022-05-04 18:02:49.040779: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] InUse at b0933e300 of size 46727168 next 17
2022-05-04 18:02:49.041233: I tensorflow/core/common_runtime/bfc_allocator.cc:1066] Free  at b0bfce300 of size 4719123712 next 18446744073709551615
2022-05-04 18:02:49.041719: I tensorflow/core/common_runtime/bfc_allocator.cc:1071]      Summary of in-use Chunks by size:
2022-05-04 18:02:49.042440: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 16 Chunks of size 256 totalling 4.0KiB
2022-05-04 18:02:49.042831: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 512 totalling 512B
2022-05-04 18:02:49.043889: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 1280 totalling 1.2KiB
2022-05-04 18:02:49.044474: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 3584 totalling 3.5KiB
2022-05-04 18:02:49.044901: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 9728 totalling 9.5KiB
2022-05-04 18:02:49.045330: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 36864 totalling 36.0KiB
2022-05-04 18:02:49.045784: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 46727168 totalling 44.56MiB
2022-05-04 18:02:49.046196: I tensorflow/core/common_runtime/bfc_allocator.cc:1078] Sum Total of in-use chunks: 44.62MiB
2022-05-04 18:02:49.046552: I tensorflow/core/common_runtime/bfc_allocator.cc:1080] total_region_allocated_bytes_: 4859428864 memory_limit_: 4859428864 available bytes: 0 curr_region_allocation_bytes_: 9718857728
2022-05-04 18:02:49.046902: I tensorflow/core/common_runtime/bfc_allocator.cc:1086] Stats:
Limit:                      4859428864
InUse:                        46783232
MaxInUse:                    140225792
NumAllocs:                          34
MaxAllocSize:                 46727168
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0

2022-05-04 18:02:49.047317: W tensorflow/core/common_runtime/bfc_allocator.cc:474] ***_________________________________________________________________________________________________
Traceback (most recent call last):
  File "D:\DatasetProcessing\ImageClassification.py", line 394, in <module>
    main()
  File "D:\DatasetProcessing\ImageClassification.py", line 387, in main
    first_model()
  File "D:\DatasetProcessing\ImageClassification.py", line 162, in first_model
    history = model.fit(train_images, train_labels, batch_size=2, epochs=4)
  File "D:\WinPython\WPy64-3980\python-3.9.8.amd64\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "D:\WinPython\WPy64-3980\python-3.9.8.amd64\lib\site-packages\tensorflow\python\framework\constant_op.py", line 102, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
tensorflow.python.framework.errors_impl.InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run _EagerConst: Dst tensor is not initialized.

I have a notebook with nvidia gtx 1060 (6GB) gpu. I've installed the newest driver available for this gpu and cuda version 11.2. I checked gpu values in task manager when script was running but it was 1% up to 5%. It looks like tensorflow is not using gpu at all.

I tried to use:

TF_GPU_ALLOCATOR=cuda_malloc_async

and

memory_limit=4096

and

allow_growth=True

I've also decreased batch_size from 128 to 2. But neither of these options worked.

Model:

model = tf.keras.Sequential([
            tf.keras.layers.Conv2D(32, (3,3), activation = 'relu', input_shape = (200, 256, 3)),
            tf.keras.layers.MaxPooling2D(2,2),
            tf.keras.layers.Conv2D(32, (3,3), activation = 'relu'),
            tf.keras.layers.MaxPooling2D(2,2),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(128, activation = tf.nn.relu),
            tf.keras.layers.Dense(20, activation = tf.nn.softmax)
        ])
model.compile(optimizer = 'adam', loss = 'sparse_categorical_crossentropy', metrics=['accuracy'])
history = model.fit(train_images, train_labels, batch_size=2, epochs=4)

It's a simple model, but I get error on model_fit function.