使用tiled对大图像执行SR时失败了

[YangFan1997]

非常感谢您的作品。

在使用参数 --tiled --tile_size 512 --tile_stride 256，并在执行前使用了

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64

对1024px*1024px的图像进行SR（V2）时，在24G显存和48G显存的机器上均可以成功执行。
但当对1536px*1536px的图像进行SR（V2）时，在24G显存和48G显存的机器上均无法执行。
错误信息如下所示，该错误信息是在48G显存的机器上出现的，它尝试申请81GB内存，注定是失败了。
在24G显存的机器上执行时，它尝试申请9GB内存，但仍然失败了。
对同一张图像处理时，申请的内存大小相差甚远，这点感觉也很奇怪。
通过代码了解到仅在Stage2时用到了tiled选项，而在Stage1时是对整个图像进行处理。
请问是否有办法解决呢
十分感谢，期待您的回复！

Traceback (most recent call last):
  File "/home/yangfan/work/DiffBIR/inference.py", line 192, in <module>
    main()
  File "/home/yangfan/work/DiffBIR/inference.py", line 187, in main
    supported_tasks[args.task](args).run()
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/yangfan/work/DiffBIR/utils/inference.py", line 156, in run
    sample = self.pipeline.run(
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/yangfan/work/DiffBIR/utils/helpers.py", line 147, in run
    clean = self.run_stage1(lq)
  File "/home/yangfan/work/DiffBIR/utils/helpers.py", line 176, in run_stage1
    clean = self.stage1_model(lq)
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/yangfan/work/DiffBIR/model/bsrnet.py", line 101, in forward
    fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/yangfan/miniconda3/envs/diffbir/lib/python3.10/site-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 81.00 GiB. GPU 0 has a total capacity of 47.53 GiB of which 19.78 GiB is free. Including non-PyTorch memory, this process has 27.54 GiB memory in use. Of the allocated memory 27.16 GiB is allocated by PyTorch, and 75.04 MiB 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)

[0x3f3f3f3fun]

感谢您的问题~，第一阶段模型确实需要同时进行tile，我最近会更新一下。