File size: 1,404 Bytes
02c5426 2e3372e 02c5426 |
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 |
from models import register
from einops import rearrange
from mmcv.cnn import build_norm_layer
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
@register('Swin_backbone')
class SwinTransformerBackbone(nn.Module):
def __init__(self, in_channels=3, embed_dims=256, depth=4, drop_path_rate=0.1):
super().__init__()
self.patch_embed = PatchEmbed(
in_channels=in_channels,
embed_dims=embed_dims,
conv_type='Conv2d',
kernel_size=5,
stride=1,
padding='same',
norm_cfg=dict(type='LN')
)
dpr = [
x.item() for x in torch.linspace(0, drop_path_rate, depth)
]
self.stage = SwinBlockSequence(
embed_dims=embed_dims,
num_heads=8,
feedforward_channels=embed_dims * 2,
depth=depth,
window_size=4,
drop_path_rate=dpr,
downsample=None
)
self.norm_layer = build_norm_layer(dict(type='LN'), embed_dims)[1]
self.out_dim = embed_dims
def forward(self, x):
x, hw_shape = self.patch_embed(x)
x, hw_shape, out, out_hw_shape = self.stage(x, hw_shape)
out = self.norm_layer(out)
x = out.view(-1, *out_hw_shape, self.out_dim).permute(0, 3, 1, 2).contiguous()
return x
|