Large Input Size to UNet model without Overlap Patching

Hi, I am trying to make a GAN model that takes an image and output an image. The model converged successfully and I can perform inference. However, my input resolution is so big (6000x8000) and I got not enough memory error. Resizing the image to lower resolution gives not good result. I tried to cut the images into smaller overlapping patches to inference, the result at each patch looked better. But when I merged the patches into the final output, there is a noticeable difference at the border and overlapping of the patches, so the image looked not whole. At present, my only solution is to feed the whole image.

I would like to ask if there is any method to inference on large input size with memory constraints, maybe saving immediate results to other storage like RAM, swap or hard disk. I can accept that the processing time will be much higher.

Any help appreciated.

Hi Duc Hoa!

You refer to “UNet model” in your title, so I will reply in that context.

In the original U-Net paper the authors have been very careful to design the
U-Net architecture so that their “overlap-tile strategy” gives the exact same
result (up to round-off error) as passing the entire image through the U-Net
all at once.

I would recommend going through that paper carefully to understand how and
why this works and then modify your U-Net, if necessary, to work the same way.

U-Net was specifically designed this way so that it can be used with arbitrarily
large images, so this is the right solution to your problem.

Yes. A U-Net saves intermediate “images” as you go down the left side of
the U so that it can recombine them with the corresponding “images” as you
come back up the right side of the U.

You can modify your U-Net code so that you copy each tensor being saved
on the right side of the U to cpu ram (assuming that you have enough of that):

x_cpu = x_gpu.cpu()
x_gpu = None   # to free x_gpu's gpu memory

and then copy it back to the gpu when you need it again on the left side of
the U. (Of course, you still need to be able to fit whatever large “image” you
are actively working on in your gpu.)

Best.

K. Frank

1 Like

Hi Frank,
Thank you so much for your detailed response. I will modify as your suggestion and review the training and data augmentation process.