Spaces:
Runtime error
Runtime error
File size: 2,182 Bytes
6ede834 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | from torch import nn
from .swin_transformer import SwinTransformer
def to_2tuple(x):
from itertools import repeat
import collections.abc
if isinstance(x, collections.abc.Iterable):
return x
return tuple(repeat(x, 2))
class ConvStem(nn.Module):
def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
super().__init__()
assert patch_size == 4
assert embed_dim % 8 == 0
img_size = to_2tuple(img_size)
patch_size = to_2tuple(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
stem = []
input_dim, output_dim = 3, embed_dim // 8
for l in range(2):
stem.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=2, padding=1, bias=False))
stem.append(nn.BatchNorm2d(output_dim))
stem.append(nn.ReLU(inplace=True))
input_dim = output_dim
output_dim *= 2
stem.append(nn.Conv2d(input_dim, embed_dim, kernel_size=1))
self.proj = nn.Sequential(*stem)
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x):
B, C, H, W = x.shape
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
x = self.proj(x)
if self.flatten:
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
x = self.norm(x)
return x
def CTransPath(
num_classes: int,
drop_rate: float = 0.,
drop_path_rate: float = 0.1,
) -> nn.Module:
model = SwinTransformer(patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24), embed_layer=ConvStem, drop_rate=drop_rate, drop_path_rate=drop_path_rate)
if num_classes == 0:
model.head = nn.Identity()
else:
model.head = nn.Linear(768, num_classes)
return model
|