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