57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
# We use the same API as https://github.com/rwightman/pytorch-image-models/blob/v0.6.11/timm/models/layers/patch_embed.py
|
|
# But we use nn.Linear instead of Conv2d and it's about 8x faster.
|
|
|
|
from functools import partial
|
|
|
|
import torch.nn as nn
|
|
from torch import _assert
|
|
from torch.nn.modules.utils import _pair
|
|
|
|
from einops import rearrange
|
|
|
|
try:
|
|
from flash_attn.ops.fused_dense import FusedDenseTD
|
|
except ImportError:
|
|
FusedDenseTD = None
|
|
|
|
|
|
class PatchEmbed(nn.Module):
|
|
""" 2D Image to Patch Embedding
|
|
"""
|
|
def __init__(
|
|
self,
|
|
img_size=224,
|
|
patch_size=16,
|
|
in_chans=3,
|
|
embed_dim=768,
|
|
norm_layer=None,
|
|
flatten=True,
|
|
bias=True,
|
|
fused_bias_fc=False,
|
|
):
|
|
super().__init__()
|
|
img_size = _pair(img_size)
|
|
patch_size = _pair(patch_size)
|
|
self.img_size = img_size
|
|
self.patch_size = patch_size
|
|
self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
|
|
self.num_patches = self.grid_size[0] * self.grid_size[1]
|
|
self.flatten = flatten
|
|
if fused_bias_fc and FusedDenseTD is None:
|
|
raise ImportError('fused_dense is not installed')
|
|
|
|
linear_cls = nn.Linear if not fused_bias_fc or not bias else FusedDenseTD
|
|
self.proj = linear_cls(in_chans * patch_size[0] * patch_size[1], embed_dim, bias=bias)
|
|
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
|
|
|
def forward(self, x):
|
|
_, _, H, W = x.shape
|
|
_assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
|
|
_assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
|
|
x = self.proj(rearrange(x, 'b c (h p1) (w p2) -> b h w (c p1 p2)',
|
|
p1=self.patch_size[0], p2=self.patch_size[1]))
|
|
if self.flatten:
|
|
x = rearrange(x, 'b h w c -> b (h w) c')
|
|
x = self.norm(x)
|
|
return x
|