Pytorch图像去噪实战（二十三）：高清大图分块去噪推理，解决显存爆炸和边缘接缝问题

Pytorch图像去噪实战（二十三）：高清大图分块去噪推理，解决显存爆炸和边缘接缝问题

一、问题场景：小图效果很好，大图直接爆显存

训练图像去噪模型时，我们通常用：

128x128 256x256

这种 patch 训练。

但真实业务里，用户上传的图片可能是：

1920x1080 3000x4000 6000x4000

我第一次直接把大图送进 UNet，结果显存直接爆了：

CUDA out of memory

即使换成 CPU，也会非常慢。

所以高清图片去噪必须做一个工程处理：

分块推理。

二、为什么大图不能直接推理？

以 UNet 为例，模型中间特征通道很多。
输入图越大，中间激活占用越大。

显存占用大致与：

H × W × 通道数 × 层数

相关。

所以图片面积增加4倍，显存可能远不止增加4倍。