import torch
import math
from torch import nn, Tensor
from torch.nn import functional as F
from typing import Optional, Callable, Union, Tuple, List
from timm.layers.helpers import to_2tuple
from timm.layers.format import Format, nchw_to
from timm.layers.trace_utils import _assert
[docs]
def resample_patch_embed(
patch_embed: nn.Parameter,
new_size: List[int],
interpolation: str = "bicubic",
antialias: bool = True,
):
"""Resample the weights of the patch embedding kernel to target resolution.
We resample the patch embedding kernel by approximately inverting the effect
of patch resizing.
Code based on:
https://github.com/google-research/big_vision/blob/b00544b81f8694488d5f36295aeb7972f3755ffe/big_vision/models/proj/flexi/vit.py
With this resizing, we can for example load a B/8 filter into a B/16 model
and, on 2x larger input image, the result will match.
Args:
patch_embed: original parameter to be resized.
new_size (tuple(int, int): target shape (height, width)-only.
interpolation (str): interpolation for resize
antialias (bool): use anti-aliasing filter in resize
verbose (bool): log operation
Returns:
Resized patch embedding kernel.
"""
import numpy as np
try:
from torch import vmap
except ImportError:
from functorch import vmap
assert len(patch_embed.shape) == 4, "Four dimensions expected"
assert len(new_size) == 2, "New shape should only be hw"
old_size = patch_embed.shape[-2:]
if tuple(old_size) == tuple(new_size):
return patch_embed
# if verbose:
# _logger.info(f"Resize patch embedding {patch_embed.shape} to {new_size}, w/ {interpolation} interpolation.")
def resize(x_np, _new_size):
x_tf = torch.Tensor(x_np)[None, None, ...]
x_upsampled = F.interpolate(
x_tf, size=_new_size, mode=interpolation, antialias=antialias
)[0, 0, ...].numpy()
return x_upsampled
def get_resize_mat(_old_size, _new_size):
mat = []
for i in range(np.prod(_old_size)):
basis_vec = np.zeros(_old_size)
basis_vec[np.unravel_index(i, _old_size)] = 1.0
mat.append(resize(basis_vec, _new_size).reshape(-1))
return np.stack(mat).T
resize_mat = get_resize_mat(old_size, new_size)
resize_mat_pinv = torch.tensor(
np.linalg.pinv(resize_mat.T), device=patch_embed.device
)
def resample_kernel(kernel):
resampled_kernel = resize_mat_pinv @ kernel.reshape(-1)
return resampled_kernel.reshape(new_size)
v_resample_kernel = vmap(vmap(resample_kernel, 0, 0), 1, 1)
orig_dtype = patch_embed.dtype
patch_embed = patch_embed.float()
patch_embed = v_resample_kernel(patch_embed)
patch_embed = patch_embed.to(orig_dtype)
return patch_embed
[docs]
class PatchEmbed(nn.Module):
"""2D Image to Patch Embedding"""
output_fmt: Format
dynamic_img_pad: torch.jit.Final[bool]
def __init__(
self,
img_size: Union[int, Tuple[int, int]] = 224,
patch_size: Union[int, Tuple[int, int]] = 16,
in_chans: int = 3,
embed_dim: int = 768,
norm_layer: Optional[Callable] = None,
flatten: bool = True,
output_fmt: Optional[str] = None,
bias: bool = True,
strict_img_size: bool = True,
dynamic_img_pad: bool = False,
):
"""
Initialize the PatchEmbed module.
Parameters
----------
img_size : int or Tuple[int, int], default=224
Input image size. If None, image size will be inferred dynamically.
patch_size : int or Tuple[int, int], default=16
Size of each image patch.
in_chans : int, default=3
Number of input channels (e.g., 3 for RGB images).
embed_dim : int, default=768
Dimension of the output patch embeddings.
norm_layer : Callable, optional
Normalization layer applied to the output embeddings.
flatten : bool, default=True
If True, flattens patches into a sequence (N, L, C).
output_fmt : str, optional
Output tensor format. If specified, overrides `flatten`.
bias : bool, default=True
Whether to include a bias term in the projection layer.
strict_img_size : bool, default=True
If True, enforces input images to match the specified size exactly.
dynamic_img_pad : bool, default=False
If True, applies dynamic padding for images not divisible by patch size.
"""
super().__init__()
self.patch_size = to_2tuple(patch_size)
self.img_size, self.grid_size, self.num_patches = self._init_img_size(img_size)
if output_fmt is not None:
self.flatten = False
self.output_fmt = Format(output_fmt)
else:
# flatten spatial dim and transpose to channels last, kept for bwd compat
self.flatten = flatten
self.output_fmt = Format.NCHW
self.strict_img_size = strict_img_size
self.dynamic_img_pad = dynamic_img_pad
self.proj = nn.Conv2d(
in_chans,
embed_dim,
kernel_size=patch_size,
stride=patch_size,
bias=bias,
)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
[docs]
def _init_img_size(self, img_size: Union[int, Tuple[int, int]]):
assert self.patch_size
if img_size is None:
return None, None, None
img_size = to_2tuple(img_size)
grid_size = tuple([s // p for s, p in zip(img_size, self.patch_size)])
num_patches = grid_size[0] * grid_size[1]
return img_size, grid_size, num_patches
[docs]
def feat_ratio(self, as_scalar=True) -> Union[Tuple[int, int], int]:
if as_scalar:
return max(self.patch_size)
else:
return self.patch_size
[docs]
def dynamic_feat_size(self, img_size: Tuple[int, int]) -> Tuple[int, int]:
"""Get grid (feature) size for given image size taking account of dynamic padding.
NOTE: must be torchscript compatible so using fixed tuple indexing
"""
if self.dynamic_img_pad:
return math.ceil(img_size[0] / self.patch_size[0]), math.ceil(
img_size[1] / self.patch_size[1]
)
else:
return (
img_size[0] // self.patch_size[0],
img_size[1] // self.patch_size[1],
)
[docs]
def forward(self, x: Tensor):
"""
Forward pass that converts an input image into patch embeddings.
Parameters
----------
x : torch.Tensor
Input tensor of shape (B, C, H, W), where
B is batch size, C is number of channels, and H, W are spatial dimensions.
Returns
-------
torch.Tensor
Patch embeddings tensor. Shape depends on output format:
- If `flatten=True`: (B, num_patches, embed_dim)
- If `flatten=False` and `output_fmt='NCHW'`: (B, embed_dim, H_p, W_p)
- If using another output format: tensor is converted accordingly.
"""
B, C, H, W = x.shape
if self.img_size is not None:
if self.strict_img_size:
_assert(
H == self.img_size[0],
f"Input height ({H}) doesn't match model ({self.img_size[0]}).",
)
_assert(
W == self.img_size[1],
f"Input width ({W}) doesn't match model ({self.img_size[1]}).",
)
elif not self.dynamic_img_pad:
_assert(
H % self.patch_size[0] == 0,
f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]}).",
)
_assert(
W % self.patch_size[1] == 0,
f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]}).",
)
if self.dynamic_img_pad:
pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
x = F.pad(x, (0, pad_w, 0, pad_h))
x = self.proj(x)
if self.flatten:
x = x.flatten(2).transpose(1, 2) # NCHW -> NLC
elif self.output_fmt != Format.NCHW:
x = nchw_to(x, self.output_fmt)
x = self.norm(x)
return x