There is a CUDA out of memory issue when running demo.py

[manadopeee]

If you set TEST_QUALITY_ENHANCEMENT = True, CUDA out of memory will occur without fail.

i'm using NVIDIA GeForce RTX 3060, memory is 12gb

Is there a way to set it up or is there a different model?

Traceback (most recent call last): File "/workspace/face-crop-plus/demo/demo.py", line 47, in cropper.process_dir(input_dir=INPUT_DIR, output_dir=OUTPUT_DIR) File "/workspace/face-crop-plus/src/face_crop_plus/cropper.py", line 909, in process_dir list(imap) File "/usr/local/lib/python3.12/dist-packages/tqdm/std.py", line 1181, in iter for obj in iterable: File "/usr/lib/python3.12/multiprocessing/pool.py", line 873, in next raise value File "/usr/lib/python3.12/multiprocessing/pool.py", line 125, in worker result = (True, func(*args, kwds)) ^^^^^^^^^^^^^^^^^^^ File "/workspace/face-crop-plus/src/face_crop_plus/cropper.py", line 836, in process_batch images = self.enh_model.predict(images, landmarks, indices) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context return func(*args, *kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/workspace/face-crop-plus/src/face_crop_plus/models/rrdb.py", line 142, in predict image_x4 = self(images[i].unsqueeze(0).div(255)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl return forward_call(*args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/workspace/face-crop-plus/src/face_crop_plus/models/rrdb.py", line 79, in forward fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2))) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, *kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl return forward_call(args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/conv.py", line 460, in forward return self._conv_forward(input, self.weight, self.bias) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/conv.py", line 456, in _conv_forward return F.conv2d(input, weight, bias, self.stride, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.00 GiB. GPU 0 has a total capacity of 11.74 GiB of which 1.64 GiB is free. Process 4028806 has 9.00 GiB memory in use. Of the allocated memory 5.77 GiB is allocated by PyTorch, and 3.10 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

[mantasu]

Would decreasing batch_size (to 6 or 4) when initializing Cropper help you?

Enhancement model is quite large but it is the only good one I found. I'll keep this open for so that I could add alternatives in the future. Thanks for pointing this out!

[manadopeee]

Thanks for your reply. I ran it with batch set to 1, but some images are still out of memory.

[mantasu]

If the detection model is used, you could try setting resize_size to something like (512, 512) which may also help