| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import xformers.ops as xops |
| | from einops import rearrange |
| | from torch.nn import functional as F |
| | import numbers |
| |
|
| | class RMSNorm(nn.Module): |
| | def __init__(self, dim, eps=1e-6): |
| | super(RMSNorm, self).__init__() |
| | self.eps = eps |
| | self.scale = nn.Parameter(torch.ones(dim)) |
| |
|
| | def forward(self, x): |
| | rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps) |
| | return self.scale * x / rms |
| |
|
| |
|
| | class ResidualBlock(nn.Module): |
| | def __init__(self, in_planes, planes, norm_fn='group', stride=1): |
| | super(ResidualBlock, self).__init__() |
| |
|
| | self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) |
| | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) |
| | self.relu = nn.ReLU(inplace=True) |
| |
|
| | num_groups = planes // 8 |
| |
|
| | if norm_fn == 'group': |
| | self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) |
| | self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) |
| | if not (stride == 1 and in_planes == planes): |
| | self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) |
| |
|
| | elif norm_fn == 'batch': |
| | self.norm1 = nn.BatchNorm2d(planes) |
| | self.norm2 = nn.BatchNorm2d(planes) |
| | if not (stride == 1 and in_planes == planes): |
| | self.norm3 = nn.BatchNorm2d(planes) |
| |
|
| | elif norm_fn == 'instance': |
| | self.norm1 = nn.InstanceNorm2d(planes) |
| | self.norm2 = nn.InstanceNorm2d(planes) |
| | if not (stride == 1 and in_planes == planes): |
| | self.norm3 = nn.InstanceNorm2d(planes) |
| |
|
| | elif norm_fn == 'none': |
| | self.norm1 = nn.Sequential() |
| | self.norm2 = nn.Sequential() |
| | if not (stride == 1 and in_planes == planes): |
| | self.norm3 = nn.Sequential() |
| |
|
| | if stride == 1 and in_planes == planes: |
| | self.downsample = None |
| |
|
| | else: |
| | self.downsample = nn.Sequential( |
| | nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3) |
| |
|
| | def forward(self, x): |
| | y = x |
| | y = self.conv1(y) |
| | y = self.norm1(y) |
| | y = self.relu(y) |
| | y = self.conv2(y) |
| | y = self.norm2(y) |
| | y = self.relu(y) |
| |
|
| | if self.downsample is not None: |
| | x = self.downsample(x) |
| |
|
| | return self.relu(x + y) |
| |
|
| |
|
| | class UnetExtractor(nn.Module): |
| | def __init__(self, in_channel=3, encoder_dim=[256, 256, 256], norm_fn='group'): |
| | super().__init__() |
| | self.in_ds = nn.Sequential( |
| | nn.Conv2d(in_channel, 64, kernel_size=7, stride=2, padding=3), |
| | nn.GroupNorm(num_groups=8, num_channels=64), |
| | nn.ReLU(inplace=True) |
| | ) |
| |
|
| | self.res1 = nn.Sequential( |
| | ResidualBlock(64, encoder_dim[0], stride=2, norm_fn=norm_fn), |
| | ResidualBlock(encoder_dim[0], encoder_dim[0], norm_fn=norm_fn) |
| | ) |
| | self.res2 = nn.Sequential( |
| | ResidualBlock(encoder_dim[0], encoder_dim[1], stride=2, norm_fn=norm_fn), |
| | ResidualBlock(encoder_dim[1], encoder_dim[1], norm_fn=norm_fn) |
| | ) |
| | self.res3 = nn.Sequential( |
| | ResidualBlock(encoder_dim[1], encoder_dim[2], stride=2, norm_fn=norm_fn), |
| | ResidualBlock(encoder_dim[2], encoder_dim[2], norm_fn=norm_fn), |
| | ) |
| |
|
| | def forward(self, x): |
| | x = self.in_ds(x) |
| | x1 = self.res1(x) |
| | x2 = self.res2(x1) |
| | x3 = self.res3(x2) |
| |
|
| | return x1, x2, x3 |
| |
|
| |
|
| | class MultiBasicEncoder(nn.Module): |
| | def __init__(self, output_dim=[128], encoder_dim=[64, 96, 128]): |
| | super(MultiBasicEncoder, self).__init__() |
| | |
| | |
| | self.conv2 = nn.Sequential( |
| | ResidualBlock(encoder_dim[2], encoder_dim[2], stride=1), |
| | nn.Conv2d(encoder_dim[2], encoder_dim[2] * 2, 3, padding=1)) |
| |
|
| | |
| | output_list = [] |
| | for dim in output_dim: |
| | conv_out = nn.Sequential( |
| | ResidualBlock(encoder_dim[2], encoder_dim[2], stride=1), |
| | nn.Conv2d(encoder_dim[2], dim[2], 3, padding=1)) |
| | output_list.append(conv_out) |
| |
|
| | self.outputs08 = nn.ModuleList(output_list) |
| |
|
| | def forward(self, x): |
| | feat1, feat2 = self.conv2(x).split(dim=0, split_size=x.shape[0] // 2) |
| |
|
| | outputs08 = [f(x) for f in self.outputs08] |
| | return outputs08, feat1, feat2 |
| | |
| |
|
| |
|
| | |
| |
|
| | def _init_weights(m): |
| | if isinstance(m, nn.Linear): |
| | nn.init.normal_(m.weight, std=.02) |
| | if m.bias is not None: |
| | nn.init.constant_(m.bias, 0) |
| |
|
| | class Mlp(nn.Module): |
| | def __init__(self, in_features, mlp_ratio=4., mlp_bias=False, out_features=None, act_layer=nn.GELU, drop=0.): |
| | super().__init__() |
| | out_features = out_features or in_features |
| | hidden_features = int(in_features * mlp_ratio) |
| | self.fc1 = nn.Linear(in_features, hidden_features, bias=mlp_bias) |
| | self.act = act_layer() |
| | self.fc2 = nn.Linear(hidden_features, out_features, bias=mlp_bias) |
| | self.drop = nn.Dropout(drop) |
| |
|
| | def forward(self, x): |
| | """ |
| | x: (B, L, D) |
| | Returns: same shape as input |
| | """ |
| | x = self.fc1(x) |
| | x = self.act(x) |
| | x = self.drop(x) |
| | x = self.fc2(x) |
| | x = self.drop(x) |
| | return x |
| |
|
| | class SelfAttention(nn.Module): |
| | def __init__(self, dim, head_dim=64, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., use_flashatt_v2=True): |
| | super().__init__() |
| | assert dim % head_dim == 0, 'dim must be divisible by head_dim' |
| | self.num_heads = dim // head_dim |
| | self.scale = qk_scale or head_dim ** -0.5 |
| |
|
| | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| | self.attn_drop_p = attn_drop |
| | self.attn_drop = nn.Dropout(attn_drop) |
| | self.proj = nn.Linear(dim, dim, bias=False) |
| | self.proj_drop = nn.Dropout(proj_drop) |
| |
|
| | self.norm_q = RMSNorm(head_dim, eps=1e-5) |
| | self.norm_k = RMSNorm(head_dim, eps=1e-5) |
| |
|
| | self.use_flashatt_v2 = use_flashatt_v2 |
| |
|
| | def forward(self, x): |
| | """ |
| | x: (B, L, D) |
| | Returns: same shape as input |
| | """ |
| | B, N, C = x.shape |
| | qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) |
| |
|
| | if self.use_flashatt_v2: |
| | qkv = qkv.permute(2, 0, 1, 3, 4) |
| | q, k, v = qkv[0], qkv[1], qkv[2] |
| | q, k = self.norm_q(q).to(v.dtype), self.norm_k(k).to(v.dtype) |
| | x = xops.memory_efficient_attention(q, k, v, op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp), p=self.attn_drop_p) |
| | x = rearrange(x, 'b n h d -> b n (h d)') |
| |
|
| | x = self.proj(x) |
| | x = self.proj_drop(x) |
| | return x |
| | |
| | class CrossAttention(nn.Module): |
| | def __init__(self, dim, head_dim=64, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., use_flashatt_v2=True): |
| | super().__init__() |
| | assert dim % head_dim == 0, 'dim must be divisible by head_dim' |
| | self.num_heads = dim // head_dim |
| | self.scale = qk_scale or head_dim ** -0.5 |
| |
|
| | self.q = nn.Linear(dim, dim, bias=qkv_bias) |
| | self.k = nn.Linear(dim, dim, bias=qkv_bias) |
| | self.v = nn.Linear(dim, dim, bias=qkv_bias) |
| | |
| | self.attn_drop_p = attn_drop |
| | self.attn_drop = nn.Dropout(attn_drop) |
| | self.proj = nn.Linear(dim, dim, bias=False) |
| | self.proj_drop = nn.Dropout(proj_drop) |
| |
|
| | self.norm_q = RMSNorm(head_dim, eps=1e-5) |
| | self.norm_k = RMSNorm(head_dim, eps=1e-5) |
| |
|
| | self.use_flashatt_v2 = use_flashatt_v2 |
| |
|
| | def forward(self, x_q, x_kv): |
| | """ |
| | x_q: query input (B, L_q, D) |
| | x_kv: key-value input (B, L_kv, D) |
| | Returns: same shape as query input (B, L_q, D) |
| | """ |
| | B, N_q, C = x_q.shape |
| | _, N_kv, _ = x_kv.shape |
| |
|
| | q = self.q(x_q).reshape(B, N_q, self.num_heads, C // self.num_heads) |
| | k = self.k(x_kv).reshape(B, N_kv, self.num_heads, C // self.num_heads) |
| | v = self.v(x_kv).reshape(B, N_kv, self.num_heads, C // self.num_heads) |
| | |
| | if self.use_flashatt_v2: |
| | q, k = self.norm_q(q).to(v.dtype), self.norm_k(k).to(v.dtype) |
| | x = xops.memory_efficient_attention( |
| | q, k, v, |
| | op=(xops.fmha.flash.FwOp, xops.fmha.flash.BwOp), |
| | p=self.attn_drop_p |
| | ) |
| | x = rearrange(x, 'b n h d -> b n (h d)') |
| |
|
| | x = self.proj(x) |
| | x = self.proj_drop(x) |
| | return x |
| |
|
| |
|
| | class TransformerBlockSelfAttn(nn.Module): |
| | def __init__(self, dim, head_dim, mlp_ratio=4., mlp_bias=False, qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_flashatt_v2=True): |
| | super().__init__() |
| | self.norm1 = norm_layer(dim, bias=False) |
| | self.attn = SelfAttention( |
| | dim, head_dim=head_dim, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, use_flashatt_v2=use_flashatt_v2) |
| | self.norm2 = norm_layer(dim, bias=False) |
| | self.mlp = Mlp(in_features=dim, mlp_ratio=mlp_ratio, mlp_bias=mlp_bias, act_layer=act_layer, drop=drop) |
| |
|
| | def forward(self, x): |
| | """ |
| | x: (B, L, D) |
| | Returns: same shape as input |
| | """ |
| | y = self.attn(self.norm1(x)) |
| | x = x + y |
| | x = x + self.mlp(self.norm2(x)) |
| | return x |
| | |
| | class TransformerBlockCrossAttn(nn.Module): |
| | def __init__(self, dim, head_dim, mlp_ratio=4., mlp_bias=False, qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_flashatt_v2=True): |
| | super().__init__() |
| | self.norm1 = norm_layer(dim, bias=False) |
| | self.attn = CrossAttention( |
| | dim, head_dim=head_dim, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, use_flashatt_v2=use_flashatt_v2) |
| | self.norm2 = norm_layer(dim, bias=False) |
| | self.mlp = Mlp(in_features=dim, mlp_ratio=mlp_ratio, mlp_bias=mlp_bias, act_layer=act_layer, drop=drop) |
| |
|
| | def forward(self, x_list): |
| | """ |
| | x_q: (B, L_q, D) |
| | x_kv: (B, L_kv, D) |
| | Returns: same shape as input |
| | """ |
| | x_q, x_kv = x_list |
| | y = self.attn(self.norm1(x_q), self.norm1(x_kv)) |
| | x = x_q + y |
| | x = x + self.mlp(self.norm2(x)) |
| | return x |
| |
|
| | class AppearanceTransformer(nn.Module): |
| | def __init__(self, num_layers, attn_dim, head_dim, ca_incides=[1, 3, 5, 7]): |
| | super().__init__() |
| | self.attn_dim = attn_dim |
| | self.num_layers = num_layers |
| | self.blocks = nn.ModuleList() |
| | self.ca_incides = ca_incides |
| |
|
| | for attn_index in range(num_layers): |
| | self.blocks.append(TransformerBlockSelfAttn(self.attn_dim, head_dim)) |
| | self.blocks[-1].apply(_init_weights) |
| |
|
| | def forward(self, x, use_checkpoint=True): |
| | """ |
| | input_tokens: (B, L, D) |
| | aggregated_tokens: List of (B, L, D) |
| | Returns: B and D remain the same, L might change if there are merge layers |
| | """ |
| | for block in self.blocks: |
| | if use_checkpoint: |
| | x = torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False) |
| | else: |
| | x = block(x) |
| |
|
| | return x |
| |
|
| |
|
| | if __name__ == '__main__': |
| | data = torch.ones((1, 3, 1024, 1024)) |
| |
|
| | model = UnetExtractor(in_channel=3, encoder_dim=[64, 96, 128]) |
| |
|
| | x1, x2, x3 = model(data) |
| | print(x1.shape, x2.shape, x3.shape) |
| |
|
