Introduction

Yes. In PyTorch, the usual way to extract image patches is to use Tensor.unfold or torch.nn.Unfold. Both expose sliding local windows over image tensors, which makes them useful for patch-based models, local feature extraction, and vision-transformer style preprocessing.

Using unfold on the Tensor

For a simple tensor-level solution, unfold works directly on the height and width dimensions.

1import torch
2
3x = torch.arange(1, 17, dtype=torch.float32).reshape(1, 1, 4, 4)
4print(x)
5
6patches = x.unfold(2, 2, 2).unfold(3, 2, 2)
7print(patches.shape)

If x has shape (batch, channels, height, width), then:

• the first unfold slices along height
• the second unfold slices along width

For this example, the result shape is:

(1, 1, 2, 2, 2, 2)

That means:

• 1 batch
• 1 channel
• 2 patch positions vertically
• 2 patch positions horizontally
• each patch is 2 x 2

Making the Patches Easier to Use

The raw unfold output is often correct but awkward. A reshape makes the patches more convenient:

patches = x.unfold(2, 2, 2).unfold(3, 2, 2)
patches = patches.contiguous().view(-1, 1, 2, 2)
print(patches)

Now each patch is an independent tensor of shape (1, 2, 2), and the total batch of patches is easier to feed into later code.

Using torch.nn.Unfold

For many model pipelines, torch.nn.Unfold is even cleaner because it behaves like an image-to-patch operator.

1import torch
2
3x = torch.arange(1, 17, dtype=torch.float32).reshape(1, 1, 4, 4)
4
5unfold = torch.nn.Unfold(kernel_size=2, stride=2)
6patches = unfold(x)
7
8print(patches.shape)
9print(patches)

This returns shape (batch, channels * kernel_height * kernel_width, number_of_patches).

For the example above, that becomes:

• batch = 1
• flattened patch size = 4
• number of patches = 4

That format is especially useful for:

• feeding local windows into linear layers
• tokenizing images for transformer-style models
• preparing patches for custom operations

Overlapping Patches

If you want overlap, use a stride smaller than the patch size.

1import torch
2
3x = torch.arange(1, 26, dtype=torch.float32).reshape(1, 1, 5, 5)
4unfold = torch.nn.Unfold(kernel_size=3, stride=1)
5patches = unfold(x)
6
7print(patches.shape)

With kernel_size=3 and stride=1, every patch overlaps heavily with its neighbors. That is common in classical image-processing pipelines and some dense prediction tasks.

Padding for Border Coverage

If the image size is not divisible by the patch geometry, you may want padding before extraction:

1import torch
2import torch.nn.functional as F
3
4x = torch.arange(1, 10, dtype=torch.float32).reshape(1, 1, 3, 3)
5x = F.pad(x, (0, 1, 0, 1))  # pad right and bottom
6
7unfold = torch.nn.Unfold(kernel_size=2, stride=2)
8patches = unfold(x)
9print(patches.shape)

Padding is often necessary when you need a complete grid of equally sized patches.

Common Pitfalls

The most common mistake is forgetting the tensor layout. PyTorch image tensors usually follow (N, C, H, W), so unfolding the wrong dimensions produces confusing shapes.

Another issue is assuming unfold returns ready-to-use image batches. It returns a structured view or flattened patch representation, so reshaping is often part of the workflow.

A third pitfall is ignoring stride. A stride equal to patch size gives non-overlapping patches, while smaller strides produce overlap. The difference changes both the patch count and the computational cost.

Finally, if you want exact border handling, think about padding explicitly. Otherwise the final rows or columns may be dropped when the dimensions do not fit the patch geometry cleanly.

Summary

• PyTorch can extract image patches with Tensor.unfold or torch.nn.Unfold.
• 'unfold is flexible and works directly on tensor dimensions.'
• 'torch.nn.Unfold is often more convenient for model pipelines.'
• Stride controls whether patches overlap.
• Padding may be needed when the image size does not divide evenly into patches.