mobilevit_xs occured CUDA out of memory?

[HTCho1]

I test my code using pretrained mobilevit_xs and pretrained HRNet(Deep High-Resolution Representation Learning for Visual Recognition) by my dataset. I checked that mobilevit_xs have less x8 parameters than HRNet and my pretrained weight of mobilevit_xs (85MB) is smaller than pretrained weight of HRNet (254MB).

When i test the model using eval image, the model using mobilevit_xs occured CUDA out of memory on RTX 3070. However, HRNet model work normally.

Why do not work the mobilevit_xs model? Eval image resolution is 1080 x 1920.

Is it because of key, query, value operations?

Error message is below. Traceback (most recent call last): File ".\streaming_points.py", line 213, in writeVideo(args) File ".\streaming_points.py", line 146, in writeVideo model = build_model(args.weight_path) File ".\streaming_points.py", line 114, in build_model _ = model(dummy_tensor) File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 224, in forward x = self.mvit0 File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 161, in forward x = self.transformer(x) File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 89, in forward x = attn(x) + x File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 31, in forward return self.fn(self.norm(x), **kwargs) File "C:\Users\user\anaconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl return forward_call(*input, **kwargs) File "C:\Users\user\anaconda3\envs\pytorch\pythonProject\Chicken Counting\live_code\Networks\MobileViT\mobilevit.py", line 70, in forward dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale RuntimeError: CUDA out of memory. Tried to allocate 1.99 GiB (GPU 0; 8.00 GiB total capacity; 3.85 GiB already allocated; 1.53 GiB free; 3.87 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_s plit_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF