DEV Community

Cover image for Upscaling a huge image on a small GPU: tiled super-resolution without seams
Wlad Radchenko
Wlad Radchenko

Posted on

Upscaling a huge image on a small GPU: tiled super-resolution without seams

You point Real-ESRGAN at a 6000x4000 photo, ask for 4x, and CUDA dies with out of memory. The model is fine. Your GPU just cannot hold a 24000x16000 tensor plus the activations to produce it. An 8GB card does not stand a chance, and even a 24GB card falls over once the input gets big enough.

The fix is to never feed the whole image to the model at once. Cut it into tiles, upscale each tile on its own, then glue the results back together. The trick is gluing them without leaving visible lines where the tiles meet. This is a walkthrough of the exact tiling code in Wunjo Make's RealESRGANer, which is the standard Real-ESRGAN inference helper (Wang et al., ICCV 2021 Workshops).

A wall of small square tiles
Tiling a picture is the same idea as tiling a wall, except the hard part is hiding the grout. Photo: Unsplash.

The two knobs that matter

Two numbers control everything here:

self.tile_size = tile      # e.g. 256: each tile is 256x256 input pixels. 0 = no tiling
self.tile_pad  = tile_pad  # default 10: extra pixels read around each tile
Enter fullscreen mode Exit fullscreen mode

tile_size is how big each chunk is in input pixels. Smaller tiles fit on smaller cards but mean more chunks and more overhead. tile_pad is the overlap, the halo of extra pixels read around each tile so the model has real neighbors to look at near the tile edge. Default is 10. Keep that number in mind, it is the whole reason seams do not show.

If tile_size is 0, the model runs on the full image and you are back to the OOM problem. Everything below assumes tile_size > 0.

Step 1. Lay out the grid

Given the padded input image, work out how many tiles fit across and down:

batch, channel, height, width = self.img.shape
output_height = height * self.scale
output_width  = width  * self.scale

# start with a black canvas at the final, upscaled size
self.output = self.img.new_zeros(output_shape)

tiles_x = math.ceil(width  / self.tile_size)
tiles_y = math.ceil(height / self.tile_size)
Enter fullscreen mode Exit fullscreen mode

Two things to notice. The output is allocated once, at full upscaled resolution, as a black image. That is the only big tensor that lives on the device for the whole run, and it is just storage, not activations, so it is cheap compared to running the model on the full frame. Each tile then writes its piece into this canvas.

The ceil means the last column and last row can be partial. A 600-pixel-wide image with tile_size=256 gives ceil(600/256) = 3 columns: 256, 256, and a runt of 88. The code handles that with a min clamp, which is the next part.

Step 2. Cut one tile, with its halo

For each tile in the grid, the code defines two rectangles. The core is the part this tile is responsible for. The padded rectangle is what actually gets read and fed to the model.

ofs_x = x * self.tile_size
ofs_y = y * self.tile_size

# core area: this tile's own region, clamped at the image edge
input_start_x = ofs_x
input_end_x   = min(ofs_x + self.tile_size, width)
input_start_y = ofs_y
input_end_y   = min(ofs_y + self.tile_size, height)

# padded area: core plus a tile_pad halo, clamped at the image edge
input_start_x_pad = max(input_start_x - self.tile_pad, 0)
input_end_x_pad   = min(input_end_x   + self.tile_pad, width)
input_start_y_pad = max(input_start_y - self.tile_pad, 0)
input_end_y_pad   = min(input_end_y   + self.tile_pad, width)  # note: clamps to width
Enter fullscreen mode Exit fullscreen mode

The max(..., 0) and min(..., width) are what keep the halo inside the image. A tile in the middle gets a full 10-pixel halo on all four sides. A tile against the left edge gets no halo on the left, because there is nothing there to read. The halo is real image data borrowed from the neighboring tile, not invented pixels.

That distinction is the entire point. A super-resolution network looks at a neighborhood around each pixel to decide what to draw. A pixel sitting one row inside a tile boundary needs to see the rows on the other side of that boundary. If you fed it a bare tile, that pixel would see the tile's edge and the model would hallucinate a border there. By reading 10 extra pixels of the actual neighbor, the model gets honest context right up to the boundary.

Then the padded tile, and only the padded tile, goes through the model:

input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad,
                            input_start_x_pad:input_end_x_pad]
with torch.no_grad():
    output_tile = self.model(input_tile)
Enter fullscreen mode Exit fullscreen mode

This is the small tensor. A 256x256 tile with a 10-pixel halo is 276x276, and at 4x its output is 1104x1104. That fits on a modest GPU even though the full output would not.

Step 3. Stitch the core back, throw the halo away

Now the careful part. The model upscaled the padded tile, halo included. But the halo pixels were also upscaled inside the neighboring tile, where they were the core. If you wrote the halo too, neighboring tiles would overwrite each other's edges and you would get exactly the seams you were trying to avoid.

So the code computes where the core sits inside the upscaled padded tile, and copies only that:

# where this tile's core lands in the final output (input coords times scale)
output_start_x = input_start_x * self.scale
output_end_x   = input_end_x   * self.scale
output_start_y = input_start_y * self.scale
output_end_y   = input_end_y   * self.scale

# where the core sits inside the upscaled padded tile
output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
output_end_x_tile   = output_start_x_tile + input_tile_width  * self.scale
output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
output_end_y_tile   = output_start_y_tile + input_tile_height * self.scale

# write only the core into the canvas
self.output[:, :, output_start_y:output_end_y, output_start_x:output_end_x] = \
    output_tile[:, :, output_start_y_tile:output_end_y_tile,
                      output_start_x_tile:output_end_x_tile]
Enter fullscreen mode Exit fullscreen mode

Read (input_start_x - input_start_x_pad) as "how many halo pixels are on the left of this tile." For an interior tile that is 10, so the slice skips the first 10*scale columns of the upscaled tile. For a left-edge tile it is 0, so nothing is skipped, which is correct because there was no left halo. Everything is multiplied by scale to move from input coordinates to output coordinates.

The result: each output pixel is written exactly once, by the tile that owns it, and every owned pixel was upscaled with real context around it. The tiles meet edge to edge with no overlap in the final image, but because each side of the boundary was computed with knowledge of the other side, the two halves line up.

A note on what "overlap" means here

This is worth being precise about, because the word overlap gets used two ways.

In this code, tiles overlap on the way in (the halo) but not on the way out. The padded read regions of two neighbors share tile_pad columns of input. The written core regions do not share anything. There is no alpha blend or feathering across the boundary. The seam is prevented by giving each side correct context, not by averaging the two sides together.

That is different from a feathered-overlap approach, where tiles overlap in the output and you cross-fade them. Both work. The halo-and-crop method here is simpler and avoids any softening from the blend, at the cost of needing the model to be consistent across the boundary, which a convolutional upscaler with enough halo generally is. With tile_pad=10, ten pixels of context is plenty for Real-ESRGAN's receptive field on its own scale.

Step 4. The whole-image padding around the edges

There is a second, smaller padding layer that is easy to miss. Before any tiling, pre_process pads the entire image:

if self.pre_pad != 0:
    self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
Enter fullscreen mode Exit fullscreen mode

This pre_pad (default 10) handles the outer border of the whole image, where there is no neighbor to borrow from. It mirrors the edge pixels outward with reflect so the model does not see a hard frame at the very edge of the picture. There is also a mod_scale pad that rounds the dimensions up to a multiple the network needs (mod 2 when scale==2, mod 4 when scale==1).

Both pads are cropped back off at the end, scaled up by the upscale factor:

# remove the mod pad, then the pre pad, both at output scale
self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale,
                                0:w - self.mod_pad_w * self.scale]
self.output = self.output[:, :, 0:h - self.pre_pad   * self.scale,
                                0:w - self.pre_pad   * self.scale]
Enter fullscreen mode Exit fullscreen mode

So you get back exactly input_size * scale, no leftover border.

Gotchas I would warn you about

A few things that bite people:

  • Pick tile_size to fit VRAM, not to be tidy. Memory scales with tile area, not perimeter. Going from 512 to 256 roughly quarters the per-tile activation memory. If you OOM, halve the tile size before touching anything else.
  • tile_pad is in input pixels, halo cost is per tile. A bigger halo means safer seams but more redundant compute, since every halo pixel is upscaled twice (once as halo, once as someone's core). The default 10 is a good balance, do not crank it to 100.
  • The OOM is swallowed. The model call is wrapped in try/except RuntimeError that just prints the error. If a single tile is still too big, you get a black square in the output instead of a crash, which can be confusing. If you see a black tile, your tile is still too large for the card.
  • Smaller tiles are slower. More tiles means more kernel launches and more halo overhead. There is a real time-versus-memory trade here. On a roomy GPU, use the biggest tile that fits.
  • Edge tiles are smaller. The last row and column are partial, and the halo is asymmetric there. The code already handles this with the min/max clamps, but if you ever rewrite this loop, that is the case that breaks first.

Wlad Radchenko About the author. I'm Wlad Radchenko, a software engineer. The code in this article comes from Wunjo Make (open source), local software for video makers, and Wunjo Design, an offline PWA for designers. Get in touch to find more on GitHub and LinkedIn.

Takeaway

Tiled super-resolution is three moves. Allocate the full output once as a black canvas. For each tile, read a core region plus a tile_pad halo of real neighbor pixels, upscale that padded tile, then write back only the core and discard the halo. The halo gives the model honest context at every boundary, the crop makes sure each output pixel is owned by exactly one tile, and the two together mean the tiles line up without a blend.

That is how an 8GB card upscales a photo whose full output would need ten times that. The model never sees more than one tile at a time.

The full method is tile_process in visual_processing/enhancement/realesrgan.py in the Wunjo Make repo. If you are hitting OOM on large inputs, start by setting a tile_size and leave tile_pad at 10.

References

  • Wang, Xie, Dong, Shan. "Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data." ICCV 2021 Workshops. arXiv:2107.10833

Top comments (0)