search

yingkaisha / keras-unet-collection

The Tensorflow, Keras implementation of U-net, V-net, U-net++, UNET 3+, Attention U-net, R2U-net, ResUnet-a, U^2-Net, TransUNET, and Swin-UNET with optional ImageNet-trained backbones.
MIT License
654 stars 182 forks source link

ResourceExhaustedError: Swin Unet #50

Open Tarandeep97 opened 2 years ago

Tarandeep97 commented 2 years ago

I have a 6 class semantic segmentation problem. I am trying to use SwinUNet as follows models.swin_unet_2d((512, 512, 3), filter_num_begin=64, n_labels=6, depth=4, stack_num_down=2, stack_num_up=2, patch_size=(2, 2), num_heads=[4, 8, 8, 8], window_size=[4, 2, 2, 2], num_mlp=512, output_activation='Softmax', shift_window=True, name='swin_unet')

But getting below error

ResourceExhaustedError: 2 root error(s) found. (0) Resource exhausted: OOM when allocating tensor with shape[8,65536,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [[node swin_unet_model/swin_transformer_block_15/name1_norm2/batchnorm/mul_2 (defined at /opt/conda/lib/python3.7/site-packages/keras_unet_collection/transformer_layers.py:623) ]] Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

 [[gradient_tape/swin_unet_model/patch_embedding_2/embedding_2/embedding_lookup/Reshape/_734]]

Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

(1) Resource exhausted: OOM when allocating tensor with shape[8,65536,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [[node swin_unet_model/swin_transformer_block_15/name1_norm2/batchnorm/mul_2 (defined at /opt/conda/lib/python3.7/site-packages/keras_unet_collection/transformer_layers.py:623) ]] Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

0 successful operations. 0 derived errors ignored. [Op:__inference_train_function_152577]

Function call stack: train_function -> train_function ` Please help.

lukestruggle commented 2 years ago

I also have the problem,don't know how to solve it