nvidia/C-RADIO
3012k
1#!/usr/bin/env python32 3# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.4#5# NVIDIA CORPORATION and its licensors retain all intellectual property6# and proprietary rights in and to this software, related documentation7# and any modifications thereto. Any use, reproduction, disclosure or8# distribution of this software and related documentation without an express9# license agreement from NVIDIA CORPORATION is strictly prohibited.10 11# E-RADIO (FasterViTv2) model from12# Mike Ranzinger, Greg Heinrich, Jan Kautz, and Pavlo Molchanov. "AM-RADIO: Agglomerative Model--Reduce All Domains Into One." arXiv preprint arXiv:2312.06709 (2023).13 14# based on FasterViT, Swin Transformer, YOLOv815 16# FasterViT:17# Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M. Alvarez, Jan Kautz, and Pavlo Molchanov. "FasterViT: Fast Vision Transformers with Hierarchical Attention." arXiv preprint arXiv:2306.06189 (2023).18 19import timm20import torch21import torch.nn as nn22from timm.models.registry import register_model23 24from timm.models.layers import trunc_normal_, DropPath, LayerNorm2d25import numpy as np26import torch.nn.functional as F27import math28import warnings29 30#######################31## Codebase from YOLOv832## BEGINNING33#######################34 35class C2f(nn.Module):36 """Faster Implementation of CSP Bottleneck with 2 convolutions."""37 """From YOLOv8 codebase"""38 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=None): # ch_in, ch_out, number, shortcut, groups, expansion39 super().__init__()40 if drop_path is None:41 drop_path = [0.0] * n42 43 self.c = int(c2 * e) # hidden channels44 self.cv1 = Conv(c1, 2 * self.c, 1, 1)45 self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2)46 self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0, drop_path=drop_path[i]) for i in range(n))47 48 def forward(self, x):49 """Forward pass through C2f layer."""50 y = list(self.cv1(x).chunk(2, 1))51 y.extend(m(y[-1]) for m in self.m)52 return self.cv2(torch.cat(y, 1))53 54 def forward_split(self, x):55 """Forward pass using split() instead of chunk()."""56 y = list(self.cv1(x).split((self.c, self.c), 1))57 y.extend(m(y[-1]) for m in self.m)58 return self.cv2(torch.cat(y, 1))59 60class Bottleneck(nn.Module):61 """Standard bottleneck."""62 63 def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, drop_path=0.0): # ch_in, ch_out, shortcut, groups, kernels, expand64 super().__init__()65 c_ = int(c2 * e) # hidden channels66 self.cv1 = Conv(c1, c_, k[0], 1)67 self.cv2 = Conv(c_, c2, k[1], 1, g=g)68 self.add = shortcut and c1 == c269 self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()70 71 def forward(self, x):72 """'forward()' applies the YOLOv5 FPN to input data."""73 return x + self.drop_path1(self.cv2(self.cv1(x))) if self.add else self.cv2(self.cv1(x))74 75 76class Conv(nn.Module):77 """Modified to support layer fusion"""78 default_act = nn.SiLU() # default activation79 80 def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, dilation=1, bn_weight_init=1, bias=False, act=True):81 super().__init__()82 83 self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, autopad(kernel_size, padding, dilation), dilation, g, bias=False)84 if 1:85 self.bn = torch.nn.BatchNorm2d(b)86 torch.nn.init.constant_(self.bn.weight, bn_weight_init)87 torch.nn.init.constant_(self.bn.bias, 0)88 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()89 90 91 def forward(self,x):92 x = self.conv(x)93 x = self.bn(x)94 x = self.act(x)95 return x96 97 @torch.no_grad()98 def switch_to_deploy(self):99 # return 1100 if not isinstance(self.bn, nn.Identity):101 c, bn = self.conv, self.bn102 w = bn.weight / (bn.running_var + bn.eps) ** 0.5103 w = c.weight * w[:, None, None, None]104 b = bn.bias - bn.running_mean * bn.weight / \105 (bn.running_var + bn.eps)**0.5106 107 self.conv.weight.data.copy_(w)108 self.conv.bias = nn.Parameter(b)109 110 self.bn = nn.Identity()111 112def autopad(k, p=None, d=1): # kernel, padding, dilation113 """Pad to 'same' shape outputs."""114 if d > 1:115 k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size116 if p is None:117 p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad118 return p119 120 121#######################122## Codebase from YOLOv8123## END124#######################125 126def pixel_unshuffle(data, factor=2):127 # performs nn.PixelShuffle(factor) in reverse, torch has some bug for ONNX and TRT, so doing it manually128 B, C, H, W = data.shape129 return data.view(B, C, factor, H//factor, factor, W//factor).permute(0,1,2,4,3,5).reshape(B, -1, H//factor, W//factor)130 131class SwiGLU(nn.Module):132 # should be more advanced, but doesnt improve results so far133 def forward(self, x):134 x, gate = x.chunk(2, dim=-1)135 return F.silu(gate) * x136 137 138def window_partition(x, window_size):139 """140 Function for partitioning image into windows and later do windowed attention141 Args:142 x: (B, C, H, W)143 window_size: window size144 Returns:145 windows - local window features (num_windows*B, window_size*window_size, C)146 (Hp, Wp) - the size of the padded image147 """148 B, C, H, W = x.shape149 150 if window_size == 0 or (window_size==H and window_size==W):151 windows = x.flatten(2).transpose(1, 2)152 Hp, Wp = H, W153 else:154 pad_h = (window_size - H % window_size) % window_size155 pad_w = (window_size - W % window_size) % window_size156 if pad_h > 0 or pad_w > 0:157 x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")158 Hp, Wp = H + pad_h, W + pad_w159 160 x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)161 windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size*window_size, C)162 163 return windows, (Hp, Wp)164 165class Conv2d_BN(nn.Module):166 '''167 Conv2d + BN layer with folding capability to speed up inference168 Can be merged with Conv() function with additional arguments169 '''170 def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bn_weight_init=1, bias=False):171 super().__init__()172 self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, padding, dilation, groups, bias=False)173 if 1:174 self.bn = torch.nn.BatchNorm2d(b)175 torch.nn.init.constant_(self.bn.weight, bn_weight_init)176 torch.nn.init.constant_(self.bn.bias, 0)177 178 def forward(self,x):179 x = self.conv(x)180 x = self.bn(x)181 return x182 183 @torch.no_grad()184 def switch_to_deploy(self):185 if not isinstance(self.bn, nn.Identity):186 c, bn = self.conv, self.bn187 w = bn.weight / (bn.running_var + bn.eps) ** 0.5188 w = c.weight * w[:, None, None, None]189 b = bn.bias - bn.running_mean * bn.weight / \190 (bn.running_var + bn.eps)**0.5191 self.conv.weight.data.copy_(w)192 self.conv.bias = nn.Parameter(b)193 self.bn = nn.Identity()194 195 196 197def window_reverse(windows, window_size, H, W, pad_hw):198 """199 Windows to the full feature map200 Args:201 windows: local window features (num_windows*B, window_size, window_size, C)202 window_size: Window size203 H: Height of image204 W: Width of image205 pad_w - a tuple of image passing used in windowing step206 Returns:207 x: (B, C, H, W)208 209 """210 # print(f"window_reverse, windows.shape {windows.shape}")211 Hp, Wp = pad_hw212 if window_size == 0 or (window_size==H and window_size==W):213 B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))214 x = windows.transpose(1, 2).view(B, -1, H, W)215 else:216 B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))217 x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)218 x = x.permute(0, 5, 1, 3, 2, 4).reshape(B,windows.shape[2], Hp, Wp)219 220 if Hp > H or Wp > W:221 x = x[:, :, :H, :W, ].contiguous()222 223 return x224 225 226 227class PosEmbMLPSwinv2D(nn.Module):228 """229 2D positional embedding from Swin Transformer v2230 Added functionality to store the positional embedding in the model and not recompute it every time231 """232 def __init__(233 self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False, cpb_mlp_hidden=512,234 ):235 super().__init__()236 self.window_size = window_size237 self.num_heads = num_heads238 # mlp to generate continuous relative position bias239 self.cpb_mlp = nn.Sequential(240 nn.Linear(2, cpb_mlp_hidden, bias=True),241 nn.ReLU(inplace=True),242 nn.Linear(cpb_mlp_hidden, num_heads, bias=False),243 )244 245 self.grid_exists = False246 self.seq_length = seq_length247 self.deploy = False248 self.num_heads = num_heads249 self.no_log = no_log250 self.pretrained_window_size = pretrained_window_size251 self.relative_bias_window_size = window_size252 253 relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(window_size, num_heads,254 pretrained_window_size, seq_length,255 no_log)256 257 self.register_buffer("relative_coords_table", relative_coords_table)258 self.register_buffer("relative_position_index", relative_position_index)259 self.register_buffer("relative_bias", relative_bias) # for EMA260 261 def relative_bias_initialization(self, window_size, num_heads, pretrained_window_size, seq_length, no_log):262 # as in separate function to support window size chage after model weights loading263 relative_coords_h = torch.arange(264 -(window_size[0] - 1), window_size[0], dtype=torch.float32265 )266 relative_coords_w = torch.arange(267 -(window_size[1] - 1), window_size[1], dtype=torch.float32268 )269 relative_coords_table = (270 torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))271 .permute(1, 2, 0)272 .contiguous()273 .unsqueeze(0)274 ) # 1, 2*Wh-1, 2*Ww-1, 2275 if pretrained_window_size[0] > 0:276 relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1277 relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1278 else:279 relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1280 relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1281 282 if not no_log:283 relative_coords_table *= 8 # normalize to -8, 8284 relative_coords_table = (285 torch.sign(relative_coords_table)286 * torch.log2(torch.abs(relative_coords_table) + 1.0)287 / np.log2(8)288 )289 290 # get pair-wise relative position index for each token inside the window291 coords_h = torch.arange(self.window_size[0])292 coords_w = torch.arange(self.window_size[1])293 coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww294 coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww295 relative_coords = (296 coords_flatten[:, :, None] - coords_flatten[:, None, :]297 ) # 2, Wh*Ww, Wh*Ww298 relative_coords = relative_coords.permute(299 1, 2, 0300 ).contiguous() # Wh*Ww, Wh*Ww, 2301 relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0302 relative_coords[:, :, 1] += self.window_size[1] - 1303 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1304 relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww305 306 relative_bias = torch.zeros(1, num_heads, seq_length, seq_length)307 308 self.relative_bias_window_size = window_size309 310 return relative_coords_table, relative_position_index, relative_bias311 312 313 def switch_to_deploy(self):314 self.deploy = True315 self.grid_exists = True316 317 def forward(self, input_tensor):318 # for efficiency, we want this forward to be folded into a single operation (sum)319 # if resolution stays the same, then we dont need to recompute MLP layers320 321 if not self.deploy or self.training:322 self.grid_exists = False323 324 #compare if all elements in self.window_size list match those in self.relative_bias_window_size325 if not all([self.window_size[i] == self.relative_bias_window_size[i] for i in range(len(self.window_size))]):326 relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(self.window_size, self.num_heads,327 self.pretrained_window_size, self.seq_length,328 self.no_log)329 330 self.relative_coords_table = relative_coords_table.to(self.relative_coords_table.device)331 self.relative_position_index = relative_position_index.to(self.relative_position_index.device)332 self.relative_bias = relative_bias.to(self.relative_bias.device)333 334 if self.deploy and self.grid_exists:335 input_tensor = input_tensor + self.relative_bias336 return input_tensor337 338 if 1:339 self.grid_exists = True340 341 relative_position_bias_table = self.cpb_mlp(342 self.relative_coords_table343 ).view(-1, self.num_heads)344 relative_position_bias = relative_position_bias_table[345 self.relative_position_index.view(-1)346 ].view(347 self.window_size[0] * self.window_size[1],348 self.window_size[0] * self.window_size[1],349 -1,350 ) # Wh*Ww,Wh*Ww,nH351 352 relative_position_bias = relative_position_bias.permute(353 2, 0, 1354 ).contiguous() # nH, Wh*Ww, Wh*Ww355 relative_position_bias = 16 * torch.sigmoid(relative_position_bias)356 357 self.relative_bias = relative_position_bias.unsqueeze(0)358 359 input_tensor = input_tensor + self.relative_bias360 return input_tensor361 362 363class GRAAttentionBlock(nn.Module):364 def __init__(self, window_size, dim_in, dim_out,365 num_heads, drop_path=0., qk_scale=None, qkv_bias=False,366 norm_layer=nn.LayerNorm, layer_scale=None,367 use_swiglu=True,368 subsample_ratio=1, dim_ratio=1, conv_base=False,369 do_windowing=True, multi_query=False, use_shift=0,370 cpb_mlp_hidden=512, conv_groups_ratio=0):371 '''372 Global Resolution Attention Block , see README for details373 Attention with subsampling to get a bigger receptive field for attention374 conv_base - use conv2d instead of avgpool2d for downsample / upsample375 376 377 '''378 super().__init__()379 380 self.shift_size=window_size//2 if use_shift else 0381 382 self.do_windowing = do_windowing383 self.subsample_ratio = subsample_ratio384 385 386 387 if do_windowing:388 if conv_base:389 self.downsample_op = nn.Conv2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()390 391 392 self.downsample_mixer = nn.Identity()393 self.upsample_mixer = nn.Identity()394 self.upsample_op = nn.ConvTranspose2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()395 else:396 self.downsample_op = nn.AvgPool2d(kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()397 self.downsample_mixer = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1) if subsample_ratio > 1 else nn.Identity()398 self.upsample_mixer = nn.Upsample(scale_factor=subsample_ratio, mode='nearest') if subsample_ratio > 1 else nn.Identity()399 self.upsample_op = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False) if subsample_ratio > 1 else nn.Identity()400 401 402 # in case there is no downsampling conv we want to have it separately403 # will help with information propagation between windows404 if subsample_ratio == 1:405 # conv_groups_ratio=0406 self.pre_conv = Conv2d_BN(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)407 # self.pre_conv = nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)408 # self.pre_conv_act = nn.ReLU6()409 #for simplicity:410 self.pre_conv_act = nn.Identity()411 if conv_groups_ratio == -1:412 self.pre_conv = nn.Identity()413 self.pre_conv_act = nn.Identity()414 415 self.window_size = window_size416 417 self.norm1 = norm_layer(dim_in)418 419 self.attn = WindowAttention(420 dim_in,421 num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,422 resolution=window_size,423 seq_length=window_size**2, dim_out=dim_in, multi_query=multi_query,424 shift_size=self.shift_size, cpb_mlp_hidden=cpb_mlp_hidden)425 426 self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()427 428 use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]429 self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim_in)) if use_layer_scale else 1430 431 ### mlp layer432 mlp_ratio = 4433 self.norm2 = norm_layer(dim_in)434 mlp_hidden_dim = int(dim_in * mlp_ratio)435 436 activation = nn.GELU if not use_swiglu else SwiGLU437 mlp_hidden_dim = int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim438 439 self.mlp = Mlp(in_features=dim_in, hidden_features=mlp_hidden_dim, act_layer=activation, use_swiglu=use_swiglu)440 441 self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1442 self.drop_path2=DropPath(drop_path) if drop_path > 0. else nn.Identity()443 444 445 def forward(self, x):446 skip_connection = x447 attn_mask = None448 449 # in case there is no downsampling conv we want to have it separately450 # will help with information propagation451 if self.subsample_ratio == 1:452 x = self.pre_conv_act(self.pre_conv(x)) + skip_connection453 454 if self.do_windowing:455 # performing windowing if required456 x = self.downsample_op(x)457 x = self.downsample_mixer(x)458 459 if self.window_size>0:460 H, W = x.shape[2], x.shape[3]461 462 if self.shift_size > 0 and H>self.window_size and W>self.window_size:463 # @swin like cyclic shift, doesnt show better performance464 x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3))465 466 x, pad_hw = window_partition(x, self.window_size)467 468 if self.shift_size > 0 and H>self.window_size and W>self.window_size:469 # set atten matrix to have -100 and the top right square470 # attn[:, :, :-self.shift_size, -self.shift_size:] = -100.0471 # calculate attention mask for SW-MSA472 # not used in final version, can be useful for some cases especially for high res473 H, W = pad_hw474 img_mask = torch.zeros((1, H, W, 1), device=x.device) # 1 H W 1475 h_slices = (slice(0, -self.window_size),476 slice(-self.window_size, -self.shift_size),477 slice(-self.shift_size, None))478 w_slices = (slice(0, -self.window_size),479 slice(-self.window_size, -self.shift_size),480 slice(-self.shift_size, None))481 cnt = 0482 for h in h_slices:483 for w in w_slices:484 img_mask[:, h, w, :] = cnt485 cnt += 1486 img_mask = img_mask.transpose(1,2).transpose(1,3)487 mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1488 489 mask_windows = mask_windows[0].view(-1, self.window_size * self.window_size)490 attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)491 attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))492 493 # window attention494 x = x + self.drop_path1(self.gamma1*self.attn(self.norm1(x), attn_mask=attn_mask)) # or pass H,W495 # mlp layer496 x = x + self.drop_path2(self.gamma2*self.mlp(self.norm2(x)))497 498 if self.do_windowing:499 if self.window_size > 0:500 x = window_reverse(x, self.window_size, H, W, pad_hw)501 502 # reverse cyclic shift503 if self.shift_size > 0 and H>self.window_size and W>self.window_size:504 # @swin like cyclic shift, not tested505 x = torch.roll(x, shifts=(self.shift_size, self.shift_size), dims=(2, 3))506 507 x = self.upsample_mixer(x)508 x = self.upsample_op(x)509 510 511 if x.shape[2] != skip_connection.shape[2] or x.shape[3] != skip_connection.shape[3]:512 x = torch.nn.functional.pad(x, ( 0, -x.shape[3] + skip_connection.shape[3], 0, -x.shape[2] + skip_connection.shape[2]), mode="reflect")513 # need to add skip connection because downsampling and upsampling will break residual connection514 # 0.5 is needed to make sure that the skip connection is not too strong515 # in case of no downsample / upsample we can show that 0.5 compensates for the residual connection516 x = 0.5 * x + 0.5 * skip_connection517 return x518 519 520 521 522class MultiResolutionAttention(nn.Module):523 """524 MultiResolutionAttention (MRA) module525 The idea is to use multiple attention blocks with different resolution526 Feature maps are downsampled / upsampled for each attention block on different blocks527 Every attention block supports windowing528 """529 530 def __init__(self, window_size, sr_ratio,531 dim, dim_ratio, num_heads,532 do_windowing=True,533 layer_scale=1e-5, norm_layer=nn.LayerNorm,534 drop_path = 0, qkv_bias=False, qk_scale=1.0,535 use_swiglu=True, multi_query=False, conv_base=False,536 use_shift=0, cpb_mlp_hidden=512, conv_groups_ratio=0) -> None:537 """538 Args:539 input_resolution: input image resolution540 window_size: window size541 compression_ratio: compression ratio542 max_depth: maximum depth of the GRA module543 use_shift: do window shifting544 """545 super().__init__()546 547 depth = len(sr_ratio)548 549 self.attention_blocks = nn.ModuleList()550 551 552 for i in range(depth):553 subsample_ratio = sr_ratio[i]554 if len(window_size) > i:555 window_size_local = window_size[i]556 else:557 window_size_local = window_size[0]558 559 self.attention_blocks.append(GRAAttentionBlock(window_size=window_size_local,560 dim_in=dim, dim_out=dim, num_heads=num_heads,561 qkv_bias=qkv_bias, qk_scale=qk_scale, norm_layer=norm_layer,562 layer_scale=layer_scale, drop_path=drop_path,563 use_swiglu=use_swiglu, subsample_ratio=subsample_ratio, dim_ratio=dim_ratio,564 do_windowing=do_windowing, multi_query=multi_query, conv_base=conv_base,565 use_shift=use_shift, cpb_mlp_hidden=cpb_mlp_hidden, conv_groups_ratio=conv_groups_ratio),566 )567 568 def forward(self, x):569 570 for attention_block in self.attention_blocks:571 x = attention_block(x)572 573 return x574 575 576 577class Mlp(nn.Module):578 """579 Multi-Layer Perceptron (MLP) block580 """581 582 def __init__(self,583 in_features,584 hidden_features=None,585 out_features=None,586 act_layer=nn.GELU,587 use_swiglu=True,588 drop=0.):589 """590 Args:591 in_features: input features dimension.592 hidden_features: hidden features dimension.593 out_features: output features dimension.594 act_layer: activation function.595 drop: dropout rate.596 """597 598 super().__init__()599 out_features = out_features or in_features600 hidden_features = hidden_features or in_features601 self.fc1 = nn.Linear(in_features, hidden_features * (2 if use_swiglu else 1), bias=False)602 self.act = act_layer()603 self.fc2 = nn.Linear(hidden_features, out_features, bias=False)604 605 def forward(self, x):606 x_size = x.size()607 x = x.view(-1, x_size[-1])608 x = self.fc1(x)609 x = self.act(x)610 x = self.fc2(x)611 x = x.view(x_size)612 return x613 614class Downsample(nn.Module):615 """616 Down-sampling block617 Pixel Unshuffle is used for down-sampling, works great accuracy - wise but takes 10% more TRT time618 """619 620 def __init__(self,621 dim,622 shuffle = False,623 ):624 """625 Args:626 dim: feature size dimension.627 shuffle: idea with628 keep_dim: bool argument for maintaining the resolution.629 """630 631 super().__init__()632 dim_out = 2 * dim633 634 if shuffle:635 self.norm = lambda x: pixel_unshuffle(x, factor=2)636 self.reduction = Conv2d_BN(dim*4, dim_out, 1, 1, 0, bias=False)637 # pixel unshuffleging works well but doesnt provide any speedup638 else:639 # removed layer norm for better, in this formulation we are getting 10% better speed640 # LayerNorm for high resolution inputs will be a pain as it pools over the entire spatial dimension641 # therefore we remove it compared to the original implementation in FasterViTv1642 self.norm = nn.Identity()643 self.reduction = Conv2d_BN(dim, dim_out, 3, 2, 1, bias=False)644 645 646 def forward(self, x):647 x = self.norm(x)648 x = self.reduction(x)649 return x650 651 652class PatchEmbed(nn.Module):653 """654 Patch embedding block655 Used to convert image into an initial set of feature maps with lower resolution656 """657 658 def __init__(self, in_chans=3, in_dim=64, dim=96, shuffle_down=False):659 """660 Args:661 in_chans: number of input channels.662 in_dim: intermediate feature size dimension to speed up stem.663 dim: final stem channel number664 shuffle_down: use PixelUnshuffle for down-sampling, effectively increases the receptive field665 """666 667 super().__init__()668 # shuffle_down = False669 if not shuffle_down:670 self.proj = nn.Identity()671 self.conv_down = nn.Sequential(672 Conv2d_BN(in_chans, in_dim, 3, 2, 1, bias=False),673 nn.ReLU(),674 Conv2d_BN(in_dim, dim, 3, 2, 1, bias=False),675 nn.ReLU()676 )677 else:678 self.proj = lambda x: pixel_unshuffle(x, factor=4)679 self.conv_down = nn.Sequential(Conv2d_BN(in_chans*16, dim, 3, 1, 1),680 nn.ReLU(),681 )682 683 def forward(self, x):684 x = self.proj(x)685 x = self.conv_down(x)686 return x687 688 689 690class ConvBlock(nn.Module):691 """692 Convolutional block, used in first couple of stages693 Experimented with plan resnet-18 like modules, they are the best in terms of throughput694 Finally, YOLOv8 idea seem to work fine (resnet-18 like block with squeezed feature dimension, and feature concatendation at the end)695 """696 def __init__(self, dim,697 drop_path=0.,698 layer_scale=None,699 kernel_size=3,700 ):701 super().__init__()702 703 self.conv1 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)704 self.act1 = nn.GELU()705 706 self.conv2 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)707 708 self.layer_scale = layer_scale709 if layer_scale is not None and type(layer_scale) in [int, float]:710 self.gamma = nn.Parameter(layer_scale * torch.ones(dim))711 self.layer_scale = True712 else:713 self.layer_scale = False714 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()715 716 def forward(self, x):717 input = x718 719 x = self.conv1(x)720 x = self.act1(x)721 x = self.conv2(x)722 723 if self.layer_scale:724 x = x * self.gamma.view(1, -1, 1, 1)725 x = input + self.drop_path(x)726 return x727 728 729class WindowAttention(nn.Module):730 # Windowed Attention from SwinV2731 # use a MLP trick to deal with various input image resolutions, then fold it to improve speed732 733 def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, resolution=0,734 seq_length=0, dim_out=None, multi_query=False, shift_size=0, cpb_mlp_hidden=512):735 # taken from EdgeViT and tweaked with attention bias.736 super().__init__()737 if not dim_out: dim_out = dim738 self.shift_size = shift_size739 self.multi_query = multi_query740 self.num_heads = num_heads741 head_dim = dim // num_heads742 self.head_dim = dim // num_heads743 744 self.dim_internal = dim745 746 self.scale = qk_scale or head_dim ** -0.5747 if not multi_query:748 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)749 else:750 self.qkv = nn.Linear(dim, dim + 2*self.head_dim, bias=qkv_bias)751 752 self.proj = nn.Linear(dim, dim_out, bias=False)753 # attention positional bias754 self.pos_emb_funct = PosEmbMLPSwinv2D(window_size=[resolution, resolution],755 pretrained_window_size=[resolution, resolution],756 num_heads=num_heads,757 seq_length=seq_length,758 cpb_mlp_hidden=cpb_mlp_hidden)759 760 self.resolution = resolution761 762 def forward(self, x, attn_mask = None):763 B, N, C = x.shape764 765 if not self.multi_query:766 qkv = self.qkv(x).reshape(B, -1, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)767 q, k, v = qkv[0], qkv[1], qkv[2]768 else:769 qkv = self.qkv(x)770 (q, k, v) = qkv.split([self.dim_internal, self.head_dim, self.head_dim], dim=2)771 772 q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)773 k = k.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)774 v = v.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)775 776 attn = (q @ k.transpose(-2, -1)) * self.scale777 778 attn = self.pos_emb_funct(attn)779 780 #add window shift781 if attn_mask is not None:782 nW = attn_mask.shape[0]783 attn = attn.view(B // nW, nW, self.num_heads, N, N) + attn_mask.unsqueeze(1).unsqueeze(0)784 attn = attn.view(-1, self.num_heads, N, N)785 786 attn = attn.softmax(dim=-1)787 x = (attn @ v).transpose(1, 2).reshape(B, -1, C)788 x = self.proj(x)789 return x790 791 792 793class FasterViTLayer(nn.Module):794 """795 fastervitlayer796 """797 798 def __init__(self,799 dim,800 depth,801 num_heads,802 window_size,803 conv=False,804 downsample=True,805 mlp_ratio=4.,806 qkv_bias=False,807 qk_scale=None,808 norm_layer=nn.LayerNorm,809 drop_path=0.,810 layer_scale=None,811 layer_scale_conv=None,812 sr_dim_ratio=1,813 sr_ratio=1,814 multi_query=False,815 use_swiglu=True,816 yolo_arch=False,817 downsample_shuffle=False,818 conv_base=False,819 use_shift=False,820 cpb_mlp_hidden=512,821 conv_groups_ratio=0,822 verbose: bool = True,823 824 ):825 """826 Args:827 dim: feature size dimension.828 depth: number of layers in each stage.829 input_resolution: input image resolution.830 window_size: window size in each stage.831 downsample: bool argument for down-sampling.832 mlp_ratio: MLP ratio.833 num_heads: number of heads in each stage.834 qkv_bias: bool argument for query, key, value learnable bias.835 qk_scale: bool argument to scaling query, key.836 drop: dropout rate.837 attn_drop: attention dropout rate.838 drop_path: drop path rate.839 norm_layer: normalization layer.840 layer_scale: layer scaling coefficient.841 use_shift: SWIN like window shifting for half the window size for every alternating layer (considering multi-resolution)842 conv_groups_ratio: group ratio for conv when no subsampling in multi-res attention843 """844 845 super().__init__()846 self.conv = conv847 self.yolo_arch=False848 self.verbose = verbose849 if conv:850 if not yolo_arch:851 self.blocks = nn.ModuleList([852 ConvBlock(dim=dim,853 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,854 layer_scale=layer_scale_conv)855 for i in range(depth)])856 self.blocks = nn.Sequential(*self.blocks)857 else:858 self.blocks = C2f(dim,dim,n=depth,shortcut=True,e=0.5)859 self.yolo_arch=True860 else:861 if not isinstance(window_size, list): window_size = [window_size]862 self.window_size = window_size[0]863 self.do_single_windowing = True864 if not isinstance(sr_ratio, list): sr_ratio = [sr_ratio]865 self.sr_ratio = sr_ratio866 if any([sr!=1 for sr in sr_ratio]) or len(set(window_size))>1:867 self.do_single_windowing = False868 do_windowing = True869 else:870 self.do_single_windowing = True871 do_windowing = False872 873 #for v2_2874 if conv_groups_ratio != -1:875 self.do_single_windowing = False876 do_windowing = True877 878 self.blocks = nn.ModuleList()879 for i in range(depth):880 self.blocks.append(881 MultiResolutionAttention(window_size=window_size,882 sr_ratio=sr_ratio,883 dim=dim,884 dim_ratio = sr_dim_ratio,885 num_heads=num_heads,886 norm_layer=norm_layer,887 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,888 layer_scale=layer_scale,889 qkv_bias=qkv_bias,890 qk_scale=qk_scale,891 use_swiglu=use_swiglu,892 do_windowing=do_windowing,893 multi_query=multi_query,894 conv_base=conv_base,895 cpb_mlp_hidden=cpb_mlp_hidden,896 use_shift =0 if ((not use_shift) or ((i) % 2 == 0)) else True ,897 conv_groups_ratio=conv_groups_ratio,898 ))899 self.blocks = nn.Sequential(*self.blocks)900 901 self.transformer = not conv902 self.downsample = None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)903 904 905 def forward(self, x):906 B, C, H, W = x.shape907 908 # do padding for transforemr909 interpolate = True910 if self.transformer and interpolate:911 # Windowed Attention will split feature map into windows with the size of window_size x window_size912 # if the resolution is not divisible by window_size, we need to interpolate the feature map913 # can be done via padding, but doing so after training hurts the model performance.914 # interpolation affects the performance as well, but not as much as padding915 if isinstance(self.window_size, list) or isinstance(self.window_size, tuple):916 current_max_window_size = max(self.window_size)917 else:918 current_max_window_size = self.window_size919 920 max_window_size = max([res_upsample*current_max_window_size for res_upsample in self.sr_ratio])921 if H % max_window_size != 0 or W % max_window_size != 0:922 new_h = int(np.ceil(H/max_window_size)*max_window_size)923 new_w = int(np.ceil(W/max_window_size)*max_window_size)924 x = F.interpolate(x, size=(new_h, new_w), mode='nearest')925 if self.verbose:926 warnings.warn(f"Choosen window size is not optimal for given resolution. Interpolation of features maps will be done and it can affect the performance. Max window size is {max_window_size}, feature map size is {H}x{W}, interpolated feature map size is {new_h}x{new_w}.")927 928 929 if self.transformer and self.do_single_windowing:930 H, W = x.shape[2], x.shape[3]931 x, pad_hw = window_partition(x, self.window_size)932 933 #run main blocks934 x = self.blocks(x)935 936 if self.transformer and self.do_single_windowing:937 x = window_reverse(x, self.window_size, H, W, pad_hw)938 939 if self.transformer and interpolate:940 #lets keep original resolution, might be not ideal, but for the upsampling tower we need to keep the expected resolution.941 x = F.interpolate(x, size=(H, W), mode='nearest')942 943 if self.downsample is None:944 return x, x945 946 return self.downsample(x), x # changing to output pre downsampled features947 948 949class InterpolateLayer(nn.Module):950 def __init__(self, size=None, scale_factor=None, mode='nearest'):951 super(InterpolateLayer, self).__init__()952 self.size = size953 self.scale_factor = scale_factor954 self.mode = mode955 956 def forward(self, x):957 return F.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode)958 959 960class HiResNeck(nn.Module):961 """962 The block is used to output dense features from all stages963 Otherwise, by default, only the last stage features are returned with FasterViTv2964 """965 def __init__(self, dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled):966 967 '''968 Hi Resolution neck to support output of high res features that are useful for dense tasks.969 depths - total number of layers in the base model970 neck_start_stage - when to start the neck, 0 - start from the first stage, 1 - start from the second stage etc.971 earlier layers result in higher resolution features at the cost of compute972 full_features_head_dim - number of channels in the dense features head973 '''974 super().__init__()975 # create feature projection layers for segmentation output976 self.neck_features_proj = nn.ModuleList()977 self.neck_start_stage = neck_start_stage978 upsample_ratio = 1979 for i in range(len(depths)):980 level_n_features_output = int(dim * 2 ** i)981 982 if self.neck_start_stage > i: continue983 984 if (upsample_ratio > 1) or full_features_head_dim!=level_n_features_output:985 feature_projection = nn.Sequential()986 if False:987 feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse988 feature_projection.add_module("dconv", nn.ConvTranspose2d(level_n_features_output,989 full_features_head_dim, kernel_size=upsample_ratio, stride=upsample_ratio))990 else:991 # B, in_channels, H, W -> B, in_channels, H*upsample_ratio, W*upsample_ratio992 # print("upsample ratio", upsample_ratio, level_n_features_output, level_n_features_output)993 feature_projection.add_module("upsample", InterpolateLayer(scale_factor=upsample_ratio, mode='nearest'))994 feature_projection.add_module("conv1", nn.Conv2d(level_n_features_output, level_n_features_output, kernel_size=3, stride=1, padding=1, groups=level_n_features_output))995 feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output))996 # B, in_channels, H*upsample_ratio, W*upsample_ratio -> B, full_features_head_dim, H*upsample_ratio, W*upsample_ratio997 feature_projection.add_module("conv2", nn.Conv2d(level_n_features_output, full_features_head_dim, kernel_size=1, stride=1, padding=0))998 else:999 feature_projection = nn.Sequential()1000 1001 self.neck_features_proj.append(feature_projection)1002 1003 if i>0 and downsample_enabled[i]:1004 upsample_ratio *= 21005 1006 def forward(self, x, il_level=-1, full_features=None):1007 if self.neck_start_stage > il_level:1008 return full_features1009 1010 if full_features is None:1011 full_features = self.neck_features_proj[il_level - self.neck_start_stage](x)1012 else:1013 #upsample torch tensor x to match full_features size, and add to full_features1014 feature_projection = self.neck_features_proj[il_level - self.neck_start_stage](x)1015 if feature_projection.shape[2] != full_features.shape[2] or feature_projection.shape[3] != full_features.shape[3]:1016 feature_projection = torch.nn.functional.pad(feature_projection, ( 0, -feature_projection.shape[3] + full_features.shape[3], 0, -feature_projection.shape[2] + full_features.shape[2]))1017 full_features = full_features + feature_projection1018 return full_features1019 1020class FasterViT(nn.Module):1021 """1022 FasterViT1023 """1024 1025 def __init__(self,1026 dim,1027 in_dim,1028 depths,1029 window_size,1030 mlp_ratio,1031 num_heads,1032 drop_path_rate=0.2,1033 in_chans=3,1034 num_classes=1000,1035 qkv_bias=False,1036 qk_scale=None,1037 layer_scale=None,1038 layer_scale_conv=None,1039 layer_norm_last=False,1040 sr_ratio = [1, 1, 1, 1],1041 max_depth = -1,1042 conv_base=False,1043 use_swiglu=False,1044 multi_query=False,1045 norm_layer=nn.LayerNorm,1046 drop_uniform=False,1047 yolo_arch=False,1048 shuffle_down=False,1049 downsample_shuffle=False,1050 return_full_features=False,1051 full_features_head_dim=128,1052 neck_start_stage=1,1053 use_neck=False,1054 use_shift=False,1055 cpb_mlp_hidden=512,1056 conv_groups_ratio=0,1057 verbose: bool = False,1058 **kwargs):1059 """1060 Args:1061 dim: feature size dimension.1062 depths: number of layers in each stage.1063 window_size: window size in each stage.1064 mlp_ratio: MLP ratio.1065 num_heads: number of heads in each stage.1066 drop_path_rate: drop path rate.1067 in_chans: number of input channels.1068 num_classes: number of classes.1069 qkv_bias: bool argument for query, key, value learnable bias.1070 qk_scale: bool argument to scaling query, key.1071 drop_rate: dropout rate.1072 attn_drop_rate: attention dropout rate.1073 norm_layer: normalization layer.1074 layer_scale: layer scaling coefficient.1075 return_full_features: output dense features as well as logits1076 full_features_head_dim: number of channels in the dense features head1077 neck_start_stage: a stage id to start full feature neck. Model has 4 stages, indix starts with 01078 for 224 resolution, the output of the stage before downsample:1079 stage 0: 56x56, stage 1: 28x28, stage 2: 14x14, stage 3: 7x71080 use_neck: even for summarization embedding use neck1081 use_shift: SWIN like window shifting but without masking attention1082 conv_groups_ratio: will be used for conv blocks where there is no multires attention,1083 if 0 then normal conv,1084 if 1 then channels are independent,1085 if -1 then no conv at all1086 1087 """1088 super().__init__()1089 1090 num_features = int(dim * 2 ** (len(depths) - 1))1091 self.num_classes = num_classes1092 self.patch_embed = PatchEmbed(in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down)1093 # set return_full_features true if we want to return full features from all stages1094 self.return_full_features = return_full_features1095 self.use_neck = use_neck1096 1097 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]1098 if drop_uniform:1099 dpr = [drop_path_rate for x in range(sum(depths))]1100 1101 if not isinstance(max_depth, list): max_depth = [max_depth] * len(depths)1102 1103 self.levels = nn.ModuleList()1104 for i in range(len(depths)):1105 conv = True if (i == 0 or i == 1) else False1106 1107 level = FasterViTLayer(dim=int(dim * 2 ** i),1108 depth=depths[i],1109 num_heads=num_heads[i],1110 window_size=window_size[i],1111 mlp_ratio=mlp_ratio,1112 qkv_bias=qkv_bias,1113 qk_scale=qk_scale,1114 conv=conv,1115 drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],1116 downsample=(i < len(depths) - 1),1117 layer_scale=layer_scale,1118 layer_scale_conv=layer_scale_conv,1119 sr_ratio=sr_ratio[i],1120 use_swiglu=use_swiglu,1121 multi_query=multi_query,1122 norm_layer=norm_layer,1123 yolo_arch=yolo_arch,1124 downsample_shuffle=downsample_shuffle,1125 conv_base=conv_base,1126 cpb_mlp_hidden=cpb_mlp_hidden,1127 use_shift=use_shift,1128 conv_groups_ratio=conv_groups_ratio,1129 verbose=verbose)1130 1131 self.levels.append(level)1132 1133 if self.return_full_features or self.use_neck:1134 #num_heads1135 downsample_enabled = [self.levels[i-1].downsample is not None for i in range(len(self.levels))]1136 self.high_res_neck = HiResNeck(dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled)1137 1138 self.switched_to_deploy = False1139 1140 self.norm = LayerNorm2d(num_features) if layer_norm_last else nn.BatchNorm2d(num_features)1141 self.avgpool = nn.AdaptiveAvgPool2d(1)1142 self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()1143 self.apply(self._init_weights)1144 1145 def _init_weights(self, m):1146 if isinstance(m, nn.Linear):1147 trunc_normal_(m.weight, std=.02)1148 if isinstance(m, nn.Linear) and m.bias is not None:1149 nn.init.constant_(m.bias, 0)1150 elif isinstance(m, nn.LayerNorm):1151 nn.init.constant_(m.bias, 0)1152 nn.init.constant_(m.weight, 1.0)1153 elif isinstance(m, LayerNorm2d):1154 nn.init.constant_(m.bias, 0)1155 nn.init.constant_(m.weight, 1.0)1156 elif isinstance(m, nn.BatchNorm2d):1157 nn.init.ones_(m.weight)1158 nn.init.zeros_(m.bias)1159 1160 @torch.jit.ignore1161 def no_weight_decay_keywords(self):1162 return {'rpb'}1163 1164 def forward_features(self, x):1165 _, _, H, W = x.shape1166 if H % 32 != 0 or W % 32 != 0:1167 raise ValueError(f"E-RADIO requires input dimensions to be divisible by 32 but got H x W: {H} x {W}")1168 x = self.patch_embed(x)1169 full_features = None1170 for il, level in enumerate(self.levels):1171 x, pre_downsample_x = level(x)1172 1173 if self.return_full_features or self.use_neck:1174 full_features = self.high_res_neck(pre_downsample_x, il, full_features)1175 1176 # x = self.norm(full_features if (self.return_full_features or self.use_neck) else x)1177 x = self.norm(x) # new version for1178 1179 if not self.return_full_features:1180 return x, None1181 1182 return x, full_features1183 1184 def forward(self, x):1185 x, full_features = self.forward_features(x)1186 1187 x = self.avgpool(x)1188 x = torch.flatten(x, 1)1189 1190 x = self.head(x)1191 if full_features is not None:1192 return x, full_features1193 return x1194 1195 def switch_to_deploy(self):1196 '''1197 A method to perform model self-compression1198 merges BN into conv layers1199 converts MLP relative positional bias into precomputed buffers1200 '''