I am using AttentionAugConv2D layer form "https://github.com/lschirmer/Attention-Augmented-Convolutional-Keras-Network" link in my code. I am working on 256x256x3 size of images. whenever I start training, I get "OOM when allocating tensor with shape[2,2,256,256,256,256] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu" issue

[patilparam-edgeneural]

2022-05-28 12:16:19.811844: W tensorflow/stream_executor/platform/default/dso_loader.cc:59] Could not load dynamic library 'cudart64_101.dll'; dlerror: cudart64_101.dll not found 2022-05-28 12:16:19.811984: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine. 2022-05-28 12:17:59.868122: W tensorflow/stream_executor/platform/default/dso_loader.cc:59] Could not load dynamic library 'nvcuda.dll'; dlerror: nvcuda.dll not found 2022-05-28 12:17:59.887137: W tensorflow/stream_executor/cuda/cuda_driver.cc:312] failed call to cuInit: UNKNOWN ERROR (303) 2022-05-28 12:17:59.890233: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:169] retrieving CUDA diagnostic information for host: DESKTOP-M8C53RA 2022-05-28 12:17:59.890303: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:176] hostname: DESKTOP-M8C53RA 2022-05-28 12:18:15.246483: 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 To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags. 2022-05-28 12:18:15.254740: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x1c4cef7ff80 initialized for platform Host (this does not guarantee that XLA will be used). Devices: 2022-05-28 12:18:15.254775: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): Host, Default Version 2022-05-30 11:02:52.095422: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 535822336 exceeds 10% of free system memory. 2022-05-30 11:02:52.065682: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 535822336 exceeds 10% of free system memory. 2022-05-30 11:02:52.095431: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 68719476736 exceeds 10% of free system memory. 2022-05-30 11:02:52.651227: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at batch_matmul_op_impl.h:730 : Resource exhausted: OOM when allocating tensor with shape[2,2,65536,65536] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu 2022-05-30 11:02:59.958471: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 536870912 exceeds 10% of free system memory. 2022-05-30 11:02:59.958471: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 536870912 exceeds 10% of free system memory. 2022-05-30 11:03:09.770279: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at tile_ops.cc:223 : Resource exhausted: OOM when allocating tensor with shape[2,2,256,256,256,256] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu 2022-05-30 11:03:09.916520: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at tile_ops.cc:223 : Resource exhausted: OOM when allocating tensor with shape[2,2,256,256,256,256] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu 2022-05-30 11:14:34.514953: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at batch_matmul_op_impl.h:730 : Resource exhausted: OOM when allocating tensor with shape[2,2,65536,65536] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu 2022-05-30 11:15:06.529323: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at tile_ops.cc:223 : Resource exhausted: OOM when allocating tensor with shape[2,2,256,256,256,256] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu 2022-05-30 11:15:06.529514: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at tile_ops.cc:223 : Resource exhausted: OOM when allocating tensor with shape[2,2,256,256,256,256] and type float on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu

[lschirmer]

For now, this code is not well optimized. The weights and parameters are not consuming much memory, but the intermediate products inside the computation are. In my experiments, I have been using it between intermediary Conv layers and residual layers with smaller spatial dimensions. So, I know that it needs improvements, and I will check it in the future.

[patilparam]

Yes. Please will update as soon as possible. it will help more in research. Thank you for information