nvidia/C-RADIOv4-H
8430k
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 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 nn22try:23 from timm.models import register_model24except ImportError:25 from timm.models.registry import register_model26 27try:28 from timm.layers import trunc_normal_, DropPath, LayerNorm2d29except ImportError:30 from timm.models.layers import trunc_normal_, DropPath, LayerNorm2d31import numpy as np32import torch.nn.functional as F33import math34import warnings35 36#######################37## Codebase from YOLOv838## BEGINNING39#######################40 41class C2f(nn.Module):42 """Faster Implementation of CSP Bottleneck with 2 convolutions."""43 """From YOLOv8 codebase"""44 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=None): # ch_in, ch_out, number, shortcut, groups, expansion45 super().__init__()46 if drop_path is None:47 drop_path = [0.0] * n48 49 self.c = int(c2 * e) # hidden channels50 self.cv1 = Conv(c1, 2 * self.c, 1, 1)51 self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2)52 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))53 54 def forward(self, x):55 """Forward pass through C2f layer."""56 y = list(self.cv1(x).chunk(2, 1))57 y.extend(m(y[-1]) for m in self.m)58 return self.cv2(torch.cat(y, 1))59 60 def forward_split(self, x):61 """Forward pass using split() instead of chunk()."""62 y = list(self.cv1(x).split((self.c, self.c), 1))63 y.extend(m(y[-1]) for m in self.m)64 return self.cv2(torch.cat(y, 1))65 66class Bottleneck(nn.Module):67 """Standard bottleneck."""68 69 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, expand70 super().__init__()71 c_ = int(c2 * e) # hidden channels72 self.cv1 = Conv(c1, c_, k[0], 1)73 self.cv2 = Conv(c_, c2, k[1], 1, g=g)74 self.add = shortcut and c1 == c275 self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()76 77 def forward(self, x):78 """'forward()' applies the YOLOv5 FPN to input data."""79 return x + self.drop_path1(self.cv2(self.cv1(x))) if self.add else self.cv2(self.cv1(x))80 81 82class Conv(nn.Module):83 """Modified to support layer fusion"""84 default_act = nn.SiLU() # default activation85 86 def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, dilation=1, bn_weight_init=1, bias=False, act=True):87 super().__init__()88 89 self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, autopad(kernel_size, padding, dilation), dilation, g, bias=False)90 if 1:91 self.bn = torch.nn.BatchNorm2d(b)92 torch.nn.init.constant_(self.bn.weight, bn_weight_init)93 torch.nn.init.constant_(self.bn.bias, 0)94 self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()95 96 97 def forward(self,x):98 x = self.conv(x)99 x = self.bn(x)100 x = self.act(x)101 return x102 103 @torch.no_grad()104 def switch_to_deploy(self):105 # return 1106 if not isinstance(self.bn, nn.Identity):107 c, bn = self.conv, self.bn108 w = bn.weight / (bn.running_var + bn.eps) ** 0.5109 w = c.weight * w[:, None, None, None]110 b = bn.bias - bn.running_mean * bn.weight / \111 (bn.running_var + bn.eps)**0.5112 113 self.conv.weight.data.copy_(w)114 self.conv.bias = nn.Parameter(b)115 116 self.bn = nn.Identity()117 118def autopad(k, p=None, d=1): # kernel, padding, dilation119 """Pad to 'same' shape outputs."""120 if d > 1:121 k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size122 if p is None:123 p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad124 return p125 126 127#######################128## Codebase from YOLOv8129## END130#######################131 132def pixel_unshuffle(data, factor=2):133 # performs nn.PixelShuffle(factor) in reverse, torch has some bug for ONNX and TRT, so doing it manually134 B, C, H, W = data.shape135 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)136 137class SwiGLU(nn.Module):138 # should be more advanced, but doesnt improve results so far139 def forward(self, x):140 x, gate = x.chunk(2, dim=-1)141 return F.silu(gate) * x142 143 144def window_partition(x, window_size):145 """146 Function for partitioning image into windows and later do windowed attention147 Args:148 x: (B, C, H, W)149 window_size: window size150 Returns:151 windows - local window features (num_windows*B, window_size*window_size, C)152 (Hp, Wp) - the size of the padded image153 """154 B, C, H, W = x.shape155 156 if window_size == 0 or (window_size==H and window_size==W):157 windows = x.flatten(2).transpose(1, 2)158 Hp, Wp = H, W159 else:160 pad_h = (window_size - H % window_size) % window_size161 pad_w = (window_size - W % window_size) % window_size162 if pad_h > 0 or pad_w > 0:163 x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")164 Hp, Wp = H + pad_h, W + pad_w165 166 x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)167 windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size*window_size, C)168 169 return windows, (Hp, Wp)170 171class Conv2d_BN(nn.Module):172 '''173 Conv2d + BN layer with folding capability to speed up inference174 Can be merged with Conv() function with additional arguments175 '''176 def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bn_weight_init=1, bias=False):177 super().__init__()178 self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, padding, dilation, groups, bias=False)179 if 1:180 self.bn = torch.nn.BatchNorm2d(b)181 torch.nn.init.constant_(self.bn.weight, bn_weight_init)182 torch.nn.init.constant_(self.bn.bias, 0)183 184 def forward(self,x):185 x = self.conv(x)186 x = self.bn(x)187 return x188 189 @torch.no_grad()190 def switch_to_deploy(self):191 if not isinstance(self.bn, nn.Identity):192 c, bn = self.conv, self.bn193 w = bn.weight / (bn.running_var + bn.eps) ** 0.5194 w = c.weight * w[:, None, None, None]195 b = bn.bias - bn.running_mean * bn.weight / \196 (bn.running_var + bn.eps)**0.5197 self.conv.weight.data.copy_(w)198 self.conv.bias = nn.Parameter(b)199 self.bn = nn.Identity()200 201 202 203def window_reverse(windows, window_size, H, W, pad_hw):204 """205 Windows to the full feature map206 Args:207 windows: local window features (num_windows*B, window_size, window_size, C)208 window_size: Window size209 H: Height of image210 W: Width of image211 pad_w - a tuple of image passing used in windowing step212 Returns:213 x: (B, C, H, W)214 215 """216 # print(f"window_reverse, windows.shape {windows.shape}")217 Hp, Wp = pad_hw218 if window_size == 0 or (window_size==H and window_size==W):219 B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))220 x = windows.transpose(1, 2).view(B, -1, H, W)221 else:222 B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))223 x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)224 x = x.permute(0, 5, 1, 3, 2, 4).reshape(B,windows.shape[2], Hp, Wp)225 226 if Hp > H or Wp > W:227 x = x[:, :, :H, :W, ].contiguous()228 229 return x230 231 232 233class PosEmbMLPSwinv2D(nn.Module):234 """235 2D positional embedding from Swin Transformer v2236 Added functionality to store the positional embedding in the model and not recompute it every time237 """238 def __init__(239 self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False, cpb_mlp_hidden=512,240 ):241 super().__init__()242 self.window_size = window_size243 self.num_heads = num_heads244 # mlp to generate continuous relative position bias245 self.cpb_mlp = nn.Sequential(246 nn.Linear(2, cpb_mlp_hidden, bias=True),247 nn.ReLU(inplace=True),248 nn.Linear(cpb_mlp_hidden, num_heads, bias=False),249 )250 251 self.grid_exists = False252 self.seq_length = seq_length253 self.deploy = False254 self.num_heads = num_heads255 self.no_log = no_log256 self.pretrained_window_size = pretrained_window_size257 self.relative_bias_window_size = window_size258 259 relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(window_size, num_heads,260 pretrained_window_size, seq_length,261 no_log)262 263 self.register_buffer("relative_coords_table", relative_coords_table)264 self.register_buffer("relative_position_index", relative_position_index)265 self.register_buffer("relative_bias", relative_bias) # for EMA266 267 def relative_bias_initialization(self, window_size, num_heads, pretrained_window_size, seq_length, no_log):268 # as in separate function to support window size chage after model weights loading269 relative_coords_h = torch.arange(270 -(window_size[0] - 1), window_size[0], dtype=torch.float32271 )272 relative_coords_w = torch.arange(273 -(window_size[1] - 1), window_size[1], dtype=torch.float32274 )275 relative_coords_table = (276 torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))277 .permute(1, 2, 0)278 .contiguous()279 .unsqueeze(0)280 ) # 1, 2*Wh-1, 2*Ww-1, 2281 if pretrained_window_size[0] > 0:282 relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1283 relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1284 else:285 relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1286 relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1287 288 if not no_log:289 relative_coords_table *= 8 # normalize to -8, 8290 relative_coords_table = (291 torch.sign(relative_coords_table)292 * torch.log2(torch.abs(relative_coords_table) + 1.0)293 / np.log2(8)294 )295 296 # get pair-wise relative position index for each token inside the window297 coords_h = torch.arange(self.window_size[0])298 coords_w = torch.arange(self.window_size[1])299 coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww300 coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww301 relative_coords = (302 coords_flatten[:, :, None] - coords_flatten[:, None, :]303 ) # 2, Wh*Ww, Wh*Ww304 relative_coords = relative_coords.permute(305 1, 2, 0306 ).contiguous() # Wh*Ww, Wh*Ww, 2307 relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0308 relative_coords[:, :, 1] += self.window_size[1] - 1309 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1310 relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww311 312 relative_bias = torch.zeros(1, num_heads, seq_length, seq_length)313 314 self.relative_bias_window_size = window_size315 316 return relative_coords_table, relative_position_index, relative_bias317 318 319 def switch_to_deploy(self):320 self.deploy = True321 self.grid_exists = True322 323 def forward(self, input_tensor):324 # for efficiency, we want this forward to be folded into a single operation (sum)325 # if resolution stays the same, then we dont need to recompute MLP layers326 327 if not self.deploy or self.training:328 self.grid_exists = False329 330 #compare if all elements in self.window_size list match those in self.relative_bias_window_size331 if not all([self.window_size[i] == self.relative_bias_window_size[i] for i in range(len(self.window_size))]):332 relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(self.window_size, self.num_heads,333 self.pretrained_window_size, self.seq_length,334 self.no_log)335 336 self.relative_coords_table = relative_coords_table.to(self.relative_coords_table.device)337 self.relative_position_index = relative_position_index.to(self.relative_position_index.device)338 self.relative_bias = relative_bias.to(self.relative_bias.device)339 340 if self.deploy and self.grid_exists:341 input_tensor = input_tensor + self.relative_bias342 return input_tensor343 344 if 1:345 self.grid_exists = True346 347 relative_position_bias_table = self.cpb_mlp(348 self.relative_coords_table349 ).view(-1, self.num_heads)350 relative_position_bias = relative_position_bias_table[351 self.relative_position_index.view(-1)352 ].view(353 self.window_size[0] * self.window_size[1],354 self.window_size[0] * self.window_size[1],355 -1,356 ) # Wh*Ww,Wh*Ww,nH357 358 relative_position_bias = relative_position_bias.permute(359 2, 0, 1360 ).contiguous() # nH, Wh*Ww, Wh*Ww361 relative_position_bias = 16 * torch.sigmoid(relative_position_bias)362 363 self.relative_bias = relative_position_bias.unsqueeze(0)364 365 input_tensor = input_tensor + self.relative_bias366 return input_tensor367 368 369class GRAAttentionBlock(nn.Module):370 def __init__(self, window_size, dim_in, dim_out,371 num_heads, drop_path=0., qk_scale=None, qkv_bias=False,372 norm_layer=nn.LayerNorm, layer_scale=None,373 use_swiglu=True,374 subsample_ratio=1, dim_ratio=1, conv_base=False,375 do_windowing=True, multi_query=False, use_shift=0,376 cpb_mlp_hidden=512, conv_groups_ratio=0):377 '''378 Global Resolution Attention Block , see README for details379 Attention with subsampling to get a bigger receptive field for attention380 conv_base - use conv2d instead of avgpool2d for downsample / upsample381 382 383 '''384 super().__init__()385 386 self.shift_size=window_size//2 if use_shift else 0387 388 self.do_windowing = do_windowing389 self.subsample_ratio = subsample_ratio390 391 392 393 if do_windowing:394 if conv_base:395 self.downsample_op = nn.Conv2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()396 397 398 self.downsample_mixer = nn.Identity()399 self.upsample_mixer = nn.Identity()400 self.upsample_op = nn.ConvTranspose2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()401 else:402 self.downsample_op = nn.AvgPool2d(kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()403 self.downsample_mixer = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1) if subsample_ratio > 1 else nn.Identity()404 self.upsample_mixer = nn.Upsample(scale_factor=subsample_ratio, mode='nearest') if subsample_ratio > 1 else nn.Identity()405 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()406 407 408 # in case there is no downsampling conv we want to have it separately409 # will help with information propagation between windows410 if subsample_ratio == 1:411 # conv_groups_ratio=0412 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)413 # 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)414 # self.pre_conv_act = nn.ReLU6()415 #for simplicity:416 self.pre_conv_act = nn.Identity()417 if conv_groups_ratio == -1:418 self.pre_conv = nn.Identity()419 self.pre_conv_act = nn.Identity()420 421 self.window_size = window_size422 423 self.norm1 = norm_layer(dim_in)424 425 self.attn = WindowAttention(426 dim_in,427 num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,428 resolution=window_size,429 seq_length=window_size**2, dim_out=dim_in, multi_query=multi_query,430 shift_size=self.shift_size, cpb_mlp_hidden=cpb_mlp_hidden)431 432 self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()433 434 use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]435 self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim_in)) if use_layer_scale else 1436 437 ### mlp layer438 mlp_ratio = 4439 self.norm2 = norm_layer(dim_in)440 mlp_hidden_dim = int(dim_in * mlp_ratio)441 442 activation = nn.GELU if not use_swiglu else SwiGLU443 mlp_hidden_dim = int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim444 445 self.mlp = Mlp(in_features=dim_in, hidden_features=mlp_hidden_dim, act_layer=activation, use_swiglu=use_swiglu)446 447 self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1448 self.drop_path2=DropPath(drop_path) if drop_path > 0. else nn.Identity()449 450 451 def forward(self, x):452 skip_connection = x453 attn_mask = None454 455 # in case there is no downsampling conv we want to have it separately456 # will help with information propagation457 if self.subsample_ratio == 1:458 x = self.pre_conv_act(self.pre_conv(x)) + skip_connection459 460 if self.do_windowing:461 # performing windowing if required462 x = self.downsample_op(x)463 x = self.downsample_mixer(x)464 465 if self.window_size>0:466 H, W = x.shape[2], x.shape[3]467 468 if self.shift_size > 0 and H>self.window_size and W>self.window_size:469 # @swin like cyclic shift, doesnt show better performance470 x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3))471 472 x, pad_hw = window_partition(x, self.window_size)473 474 if self.shift_size > 0 and H>self.window_size and W>self.window_size:475 # set atten matrix to have -100 and the top right square476 # attn[:, :, :-self.shift_size, -self.shift_size:] = -100.0477 # calculate attention mask for SW-MSA478 # not used in final version, can be useful for some cases especially for high res479 H, W = pad_hw480 img_mask = torch.zeros((1, H, W, 1), device=x.device) # 1 H W 1481 h_slices = (slice(0, -self.window_size),482 slice(-self.window_size, -self.shift_size),483 slice(-self.shift_size, None))484 w_slices = (slice(0, -self.window_size),485 slice(-self.window_size, -self.shift_size),486 slice(-self.shift_size, None))487 cnt = 0488 for h in h_slices:489 for w in w_slices:490 img_mask[:, h, w, :] = cnt491 cnt += 1492 img_mask = img_mask.transpose(1,2).transpose(1,3)493 mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1494 495 mask_windows = mask_windows[0].view(-1, self.window_size * self.window_size)496 attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)497 attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))498 499 # window attention500 x = x + self.drop_path1(self.gamma1*self.attn(self.norm1(x), attn_mask=attn_mask)) # or pass H,W501 # mlp layer502 x = x + self.drop_path2(self.gamma2*self.mlp(self.norm2(x)))503 504 if self.do_windowing:505 if self.window_size > 0:506 x = window_reverse(x, self.window_size, H, W, pad_hw)507 508 # reverse cyclic shift509 if self.shift_size > 0 and H>self.window_size and W>self.window_size:510 # @swin like cyclic shift, not tested511 x = torch.roll(x, shifts=(self.shift_size, self.shift_size), dims=(2, 3))512 513 x = self.upsample_mixer(x)514 x = self.upsample_op(x)515 516 517 if x.shape[2] != skip_connection.shape[2] or x.shape[3] != skip_connection.shape[3]:518 x = torch.nn.functional.pad(x, ( 0, -x.shape[3] + skip_connection.shape[3], 0, -x.shape[2] + skip_connection.shape[2]), mode="reflect")519 # need to add skip connection because downsampling and upsampling will break residual connection520 # 0.5 is needed to make sure that the skip connection is not too strong521 # in case of no downsample / upsample we can show that 0.5 compensates for the residual connection522 x = 0.5 * x + 0.5 * skip_connection523 return x524 525 526 527 528class MultiResolutionAttention(nn.Module):529 """530 MultiResolutionAttention (MRA) module531 The idea is to use multiple attention blocks with different resolution532 Feature maps are downsampled / upsampled for each attention block on different blocks533 Every attention block supports windowing534 """535 536 def __init__(self, window_size, sr_ratio,537 dim, dim_ratio, num_heads,538 do_windowing=True,539 layer_scale=1e-5, norm_layer=nn.LayerNorm,540 drop_path = 0, qkv_bias=False, qk_scale=1.0,541 use_swiglu=True, multi_query=False, conv_base=False,542 use_shift=0, cpb_mlp_hidden=512, conv_groups_ratio=0) -> None:543 """544 Args:545 input_resolution: input image resolution546 window_size: window size547 compression_ratio: compression ratio548 max_depth: maximum depth of the GRA module549 use_shift: do window shifting550 """551 super().__init__()552 553 depth = len(sr_ratio)554 555 self.attention_blocks = nn.ModuleList()556 557 558 for i in range(depth):559 subsample_ratio = sr_ratio[i]560 if len(window_size) > i:561 window_size_local = window_size[i]562 else:563 window_size_local = window_size[0]564 565 self.attention_blocks.append(GRAAttentionBlock(window_size=window_size_local,566 dim_in=dim, dim_out=dim, num_heads=num_heads,567 qkv_bias=qkv_bias, qk_scale=qk_scale, norm_layer=norm_layer,568 layer_scale=layer_scale, drop_path=drop_path,569 use_swiglu=use_swiglu, subsample_ratio=subsample_ratio, dim_ratio=dim_ratio,570 do_windowing=do_windowing, multi_query=multi_query, conv_base=conv_base,571 use_shift=use_shift, cpb_mlp_hidden=cpb_mlp_hidden, conv_groups_ratio=conv_groups_ratio),572 )573 574 def forward(self, x):575 576 for attention_block in self.attention_blocks:577 x = attention_block(x)578 579 return x580 581 582 583class Mlp(nn.Module):584 """585 Multi-Layer Perceptron (MLP) block586 """587 588 def __init__(self,589 in_features,590 hidden_features=None,591 out_features=None,592 act_layer=nn.GELU,593 use_swiglu=True,594 drop=0.):595 """596 Args:597 in_features: input features dimension.598 hidden_features: hidden features dimension.599 out_features: output features dimension.600 act_layer: activation function.601 drop: dropout rate.602 """603 604 super().__init__()605 out_features = out_features or in_features606 hidden_features = hidden_features or in_features607 self.fc1 = nn.Linear(in_features, hidden_features * (2 if use_swiglu else 1), bias=False)608 self.act = act_layer()609 self.fc2 = nn.Linear(hidden_features, out_features, bias=False)610 611 def forward(self, x):612 x_size = x.size()613 x = x.view(-1, x_size[-1])614 x = self.fc1(x)615 x = self.act(x)616 x = self.fc2(x)617 x = x.view(x_size)618 return x619 620class Downsample(nn.Module):621 """622 Down-sampling block623 Pixel Unshuffle is used for down-sampling, works great accuracy - wise but takes 10% more TRT time624 """625 626 def __init__(self,627 dim,628 shuffle = False,629 ):630 """631 Args:632 dim: feature size dimension.633 shuffle: idea with634 keep_dim: bool argument for maintaining the resolution.635 """636 637 super().__init__()638 dim_out = 2 * dim639 640 if shuffle:641 self.norm = lambda x: pixel_unshuffle(x, factor=2)642 self.reduction = Conv2d_BN(dim*4, dim_out, 1, 1, 0, bias=False)643 # pixel unshuffleging works well but doesnt provide any speedup644 else:645 # removed layer norm for better, in this formulation we are getting 10% better speed646 # LayerNorm for high resolution inputs will be a pain as it pools over the entire spatial dimension647 # therefore we remove it compared to the original implementation in FasterViT648 self.norm = nn.Identity()649 self.reduction = Conv2d_BN(dim, dim_out, 3, 2, 1, bias=False)650 651 652 def forward(self, x):653 x = self.norm(x)654 x = self.reduction(x)655 return x656 657 658class PatchEmbed(nn.Module):659 """660 Patch embedding block661 Used to convert image into an initial set of feature maps with lower resolution662 """663 664 def __init__(self, in_chans=3, in_dim=64, dim=96, shuffle_down=False):665 """666 Args:667 in_chans: number of input channels.668 in_dim: intermediate feature size dimension to speed up stem.669 dim: final stem channel number670 shuffle_down: use PixelUnshuffle for down-sampling, effectively increases the receptive field671 """672 673 super().__init__()674 # shuffle_down = False675 if not shuffle_down:676 self.proj = nn.Identity()677 self.conv_down = nn.Sequential(678 Conv2d_BN(in_chans, in_dim, 3, 2, 1, bias=False),679 nn.ReLU(),680 Conv2d_BN(in_dim, dim, 3, 2, 1, bias=False),681 nn.ReLU()682 )683 else:684 self.proj = lambda x: pixel_unshuffle(x, factor=4)685 self.conv_down = nn.Sequential(Conv2d_BN(in_chans*16, dim, 3, 1, 1),686 nn.ReLU(),687 )688 689 def forward(self, x):690 x = self.proj(x)691 x = self.conv_down(x)692 return x693 694 695 696class ConvBlock(nn.Module):697 """698 Convolutional block, used in first couple of stages699 Experimented with plan resnet-18 like modules, they are the best in terms of throughput700 Finally, YOLOv8 idea seem to work fine (resnet-18 like block with squeezed feature dimension, and feature concatendation at the end)701 """702 def __init__(self, dim,703 drop_path=0.,704 layer_scale=None,705 kernel_size=3,706 ):707 super().__init__()708 709 self.conv1 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)710 self.act1 = nn.GELU()711 712 self.conv2 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)713 714 self.layer_scale = layer_scale715 if layer_scale is not None and type(layer_scale) in [int, float]:716 self.gamma = nn.Parameter(layer_scale * torch.ones(dim))717 self.layer_scale = True718 else:719 self.layer_scale = False720 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()721 722 def forward(self, x):723 input = x724 725 x = self.conv1(x)726 x = self.act1(x)727 x = self.conv2(x)728 729 if self.layer_scale:730 x = x * self.gamma.view(1, -1, 1, 1)731 x = input + self.drop_path(x)732 return x733 734 735class WindowAttention(nn.Module):736 # Windowed Attention from SwinV2737 # use a MLP trick to deal with various input image resolutions, then fold it to improve speed738 739 def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, resolution=0,740 seq_length=0, dim_out=None, multi_query=False, shift_size=0, cpb_mlp_hidden=512):741 # taken from EdgeViT and tweaked with attention bias.742 super().__init__()743 if not dim_out: dim_out = dim744 self.shift_size = shift_size745 self.multi_query = multi_query746 self.num_heads = num_heads747 head_dim = dim // num_heads748 self.head_dim = dim // num_heads749 750 self.dim_internal = dim751 752 self.scale = qk_scale or head_dim ** -0.5753 if not multi_query:754 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)755 else:756 self.qkv = nn.Linear(dim, dim + 2*self.head_dim, bias=qkv_bias)757 758 self.proj = nn.Linear(dim, dim_out, bias=False)759 # attention positional bias760 self.pos_emb_funct = PosEmbMLPSwinv2D(window_size=[resolution, resolution],761 pretrained_window_size=[resolution, resolution],762 num_heads=num_heads,763 seq_length=seq_length,764 cpb_mlp_hidden=cpb_mlp_hidden)765 766 self.resolution = resolution767 768 def forward(self, x, attn_mask = None):769 B, N, C = x.shape770 771 if not self.multi_query:772 qkv = self.qkv(x).reshape(B, -1, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)773 q, k, v = qkv[0], qkv[1], qkv[2]774 else:775 qkv = self.qkv(x)776 (q, k, v) = qkv.split([self.dim_internal, self.head_dim, self.head_dim], dim=2)777 778 q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)779 k = k.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)780 v = v.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)781 782 attn = (q @ k.transpose(-2, -1)) * self.scale783 784 attn = self.pos_emb_funct(attn)785 786 #add window shift787 if attn_mask is not None:788 nW = attn_mask.shape[0]789 attn = attn.view(B // nW, nW, self.num_heads, N, N) + attn_mask.unsqueeze(1).unsqueeze(0)790 attn = attn.view(-1, self.num_heads, N, N)791 792 attn = attn.softmax(dim=-1)793 x = (attn @ v).transpose(1, 2).reshape(B, -1, C)794 x = self.proj(x)795 return x796 797 798 799class ERADIOLayer(nn.Module):800 """801 E-RADIO Layer802 """803 804 def __init__(self,805 dim,806 depth,807 num_heads,808 window_size,809 conv=False,810 downsample=True,811 mlp_ratio=4.,812 qkv_bias=False,813 qk_scale=None,814 norm_layer=nn.LayerNorm,815 drop_path=0.,816 layer_scale=None,817 layer_scale_conv=None,818 sr_dim_ratio=1,819 sr_ratio=1,820 multi_query=False,821 use_swiglu=True,822 yolo_arch=False,823 downsample_shuffle=False,824 conv_base=False,825 use_shift=False,826 cpb_mlp_hidden=512,827 conv_groups_ratio=0,828 verbose: bool = True,829 830 ):831 """832 Args:833 dim: feature size dimension.834 depth: number of layers in each stage.835 input_resolution: input image resolution.836 window_size: window size in each stage.837 downsample: bool argument for down-sampling.838 mlp_ratio: MLP ratio.839 num_heads: number of heads in each stage.840 qkv_bias: bool argument for query, key, value learnable bias.841 qk_scale: bool argument to scaling query, key.842 drop: dropout rate.843 attn_drop: attention dropout rate.844 drop_path: drop path rate.845 norm_layer: normalization layer.846 layer_scale: layer scaling coefficient.847 use_shift: SWIN like window shifting for half the window size for every alternating layer (considering multi-resolution)848 conv_groups_ratio: group ratio for conv when no subsampling in multi-res attention849 """850 851 super().__init__()852 self.conv = conv853 self.yolo_arch=False854 self.verbose = verbose855 if conv:856 if not yolo_arch:857 self.blocks = nn.ModuleList([858 ConvBlock(dim=dim,859 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,860 layer_scale=layer_scale_conv)861 for i in range(depth)])862 self.blocks = nn.Sequential(*self.blocks)863 else:864 self.blocks = C2f(dim,dim,n=depth,shortcut=True,e=0.5)865 self.yolo_arch=True866 else:867 if not isinstance(window_size, list): window_size = [window_size]868 self.window_size = window_size[0]869 self.do_single_windowing = True870 if not isinstance(sr_ratio, list): sr_ratio = [sr_ratio]871 self.sr_ratio = sr_ratio872 if any([sr!=1 for sr in sr_ratio]) or len(set(window_size))>1:873 self.do_single_windowing = False874 do_windowing = True875 else:876 self.do_single_windowing = True877 do_windowing = False878 879 #for v2_2880 if conv_groups_ratio != -1:881 self.do_single_windowing = False882 do_windowing = True883 884 self.blocks = nn.ModuleList()885 for i in range(depth):886 self.blocks.append(887 MultiResolutionAttention(window_size=window_size,888 sr_ratio=sr_ratio,889 dim=dim,890 dim_ratio = sr_dim_ratio,891 num_heads=num_heads,892 norm_layer=norm_layer,893 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,894 layer_scale=layer_scale,895 qkv_bias=qkv_bias,896 qk_scale=qk_scale,897 use_swiglu=use_swiglu,898 do_windowing=do_windowing,899 multi_query=multi_query,900 conv_base=conv_base,901 cpb_mlp_hidden=cpb_mlp_hidden,902 use_shift =0 if ((not use_shift) or ((i) % 2 == 0)) else True ,903 conv_groups_ratio=conv_groups_ratio,904 ))905 self.blocks = nn.Sequential(*self.blocks)906 907 self.transformer = not conv908 self.downsample = None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)909 910 911 def forward(self, x):912 B, C, H, W = x.shape913 914 # do padding for transforemr915 interpolate = True916 if self.transformer and interpolate:917 # Windowed Attention will split feature map into windows with the size of window_size x window_size918 # if the resolution is not divisible by window_size, we need to interpolate the feature map919 # can be done via padding, but doing so after training hurts the model performance.920 # interpolation affects the performance as well, but not as much as padding921 if isinstance(self.window_size, list) or isinstance(self.window_size, tuple):922 current_max_window_size = max(self.window_size)923 else:924 current_max_window_size = self.window_size925 926 max_window_size = max([res_upsample*current_max_window_size for res_upsample in self.sr_ratio])927 if H % max_window_size != 0 or W % max_window_size != 0:928 new_h = int(np.ceil(H/max_window_size)*max_window_size)929 new_w = int(np.ceil(W/max_window_size)*max_window_size)930 x = F.interpolate(x, size=(new_h, new_w), mode='nearest')931 if self.verbose:932 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}.")933 934 935 if self.transformer and self.do_single_windowing:936 H, W = x.shape[2], x.shape[3]937 x, pad_hw = window_partition(x, self.window_size)938 939 #run main blocks940 x = self.blocks(x)941 942 if self.transformer and self.do_single_windowing:943 x = window_reverse(x, self.window_size, H, W, pad_hw)944 945 if self.transformer and interpolate:946 #lets keep original resolution, might be not ideal, but for the upsampling tower we need to keep the expected resolution.947 x = F.interpolate(x, size=(H, W), mode='nearest')948 949 if self.downsample is None:950 return x, x951 952 return self.downsample(x), x # changing to output pre downsampled features953 954 955class InterpolateLayer(nn.Module):956 def __init__(self, size=None, scale_factor=None, mode='nearest'):957 super(InterpolateLayer, self).__init__()958 self.size = size959 self.scale_factor = scale_factor960 self.mode = mode961 962 def forward(self, x):963 return F.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode)964 965 966class HiResNeck(nn.Module):967 """968 The block is used to output dense features from all stages969 Otherwise, by default, only the last stage features are returned with E-RADIO970 """971 def __init__(self, dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled):972 973 '''974 Hi Resolution neck to support output of high res features that are useful for dense tasks.975 depths - total number of layers in the base model976 neck_start_stage - when to start the neck, 0 - start from the first stage, 1 - start from the second stage etc.977 earlier layers result in higher resolution features at the cost of compute978 full_features_head_dim - number of channels in the dense features head979 '''980 super().__init__()981 # create feature projection layers for segmentation output982 self.neck_features_proj = nn.ModuleList()983 self.neck_start_stage = neck_start_stage984 upsample_ratio = 1985 for i in range(len(depths)):986 level_n_features_output = int(dim * 2 ** i)987 988 if self.neck_start_stage > i: continue989 990 if (upsample_ratio > 1) or full_features_head_dim!=level_n_features_output:991 feature_projection = nn.Sequential()992 if False:993 feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse994 feature_projection.add_module("dconv", nn.ConvTranspose2d(level_n_features_output,995 full_features_head_dim, kernel_size=upsample_ratio, stride=upsample_ratio))996 else:997 # B, in_channels, H, W -> B, in_channels, H*upsample_ratio, W*upsample_ratio998 # print("upsample ratio", upsample_ratio, level_n_features_output, level_n_features_output)999 feature_projection.add_module("upsample", InterpolateLayer(scale_factor=upsample_ratio, mode='nearest'))1000 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))1001 feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output))1002 # B, in_channels, H*upsample_ratio, W*upsample_ratio -> B, full_features_head_dim, H*upsample_ratio, W*upsample_ratio1003 feature_projection.add_module("conv2", nn.Conv2d(level_n_features_output, full_features_head_dim, kernel_size=1, stride=1, padding=0))1004 else:1005 feature_projection = nn.Sequential()1006 1007 self.neck_features_proj.append(feature_projection)1008 1009 if i>0 and downsample_enabled[i]:1010 upsample_ratio *= 21011 1012 def forward(self, x, il_level=-1, full_features=None):1013 if self.neck_start_stage > il_level:1014 return full_features1015 1016 if full_features is None:1017 full_features = self.neck_features_proj[il_level - self.neck_start_stage](x)1018 else:1019 #upsample torch tensor x to match full_features size, and add to full_features1020 feature_projection = self.neck_features_proj[il_level - self.neck_start_stage](x)1021 if feature_projection.shape[2] != full_features.shape[2] or feature_projection.shape[3] != full_features.shape[3]:1022 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]))1023 full_features = full_features + feature_projection1024 return full_features1025 1026class ERADIO(nn.Module):1027 """1028 Efficient RADIO1029 """1030 1031 def __init__(self,1032 dim,1033 in_dim,1034 depths,1035 window_size,1036 mlp_ratio,1037 num_heads,1038 drop_path_rate=0.2,1039 in_chans=3,1040 num_classes=1000,1041 qkv_bias=False,1042 qk_scale=None,1043 layer_scale=None,1044 layer_scale_conv=None,1045 layer_norm_last=False,1046 sr_ratio = [1, 1, 1, 1],1047 max_depth = -1,1048 conv_base=False,1049 use_swiglu=False,1050 multi_query=False,1051 norm_layer=nn.LayerNorm,1052 drop_uniform=False,1053 yolo_arch=False,1054 shuffle_down=False,1055 downsample_shuffle=False,1056 return_full_features=False,1057 full_features_head_dim=128,1058 neck_start_stage=1,1059 use_neck=False,1060 use_shift=False,1061 cpb_mlp_hidden=512,1062 conv_groups_ratio=0,1063 verbose: bool = False,1064 **kwargs):1065 """1066 Args:1067 dim: feature size dimension.1068 depths: number of layers in each stage.1069 window_size: window size in each stage.1070 mlp_ratio: MLP ratio.1071 num_heads: number of heads in each stage.1072 drop_path_rate: drop path rate.1073 in_chans: number of input channels.1074 num_classes: number of classes.1075 qkv_bias: bool argument for query, key, value learnable bias.1076 qk_scale: bool argument to scaling query, key.1077 drop_rate: dropout rate.1078 attn_drop_rate: attention dropout rate.1079 norm_layer: normalization layer.1080 layer_scale: layer scaling coefficient.1081 return_full_features: output dense features as well as logits1082 full_features_head_dim: number of channels in the dense features head1083 neck_start_stage: a stage id to start full feature neck. Model has 4 stages, indix starts with 01084 for 224 resolution, the output of the stage before downsample:1085 stage 0: 56x56, stage 1: 28x28, stage 2: 14x14, stage 3: 7x71086 use_neck: even for summarization embedding use neck1087 use_shift: SWIN like window shifting but without masking attention1088 conv_groups_ratio: will be used for conv blocks where there is no multires attention,1089 if 0 then normal conv,1090 if 1 then channels are independent,1091 if -1 then no conv at all1092 1093 """1094 super().__init__()1095 1096 num_features = int(dim * 2 ** (len(depths) - 1))1097 self.num_classes = num_classes1098 self.patch_embed = PatchEmbed(in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down)1099 # set return_full_features true if we want to return full features from all stages1100 self.return_full_features = return_full_features1101 self.use_neck = use_neck1102 1103 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]1104 if drop_uniform:1105 dpr = [drop_path_rate for x in range(sum(depths))]1106 1107 if not isinstance(max_depth, list): max_depth = [max_depth] * len(depths)1108 1109 self.levels = nn.ModuleList()1110 for i in range(len(depths)):1111 conv = True if (i == 0 or i == 1) else False1112 1113 level = ERADIOLayer(dim=int(dim * 2 ** i),1114 depth=depths[i],1115 num_heads=num_heads[i],1116 window_size=window_size[i],1117 mlp_ratio=mlp_ratio,1118 qkv_bias=qkv_bias,1119 qk_scale=qk_scale,1120 conv=conv,1121 drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],1122 downsample=(i < len(depths) - 1),1123 layer_scale=layer_scale,1124 layer_scale_conv=layer_scale_conv,1125 sr_ratio=sr_ratio[i],1126 use_swiglu=use_swiglu,1127 multi_query=multi_query,1128 norm_layer=norm_layer,1129 yolo_arch=yolo_arch,1130 downsample_shuffle=downsample_shuffle,1131 conv_base=conv_base,1132 cpb_mlp_hidden=cpb_mlp_hidden,1133 use_shift=use_shift,1134 conv_groups_ratio=conv_groups_ratio,1135 verbose=verbose)1136 1137 self.levels.append(level)1138 1139 if self.return_full_features or self.use_neck:1140 #num_heads1141 downsample_enabled = [self.levels[i-1].downsample is not None for i in range(len(self.levels))]1142 self.high_res_neck = HiResNeck(dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled)1143 1144 self.switched_to_deploy = False1145 1146 self.norm = LayerNorm2d(num_features) if layer_norm_last else nn.BatchNorm2d(num_features)1147 self.avgpool = nn.AdaptiveAvgPool2d(1)1148 self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()1149 self.apply(self._init_weights)1150 1151 def _init_weights(self, m):1152 if isinstance(m, nn.Linear):1153 trunc_normal_(m.weight, std=.02)1154 if isinstance(m, nn.Linear) and m.bias is not None:1155 nn.init.constant_(m.bias, 0)1156 elif isinstance(m, nn.LayerNorm):1157 nn.init.constant_(m.bias, 0)1158 nn.init.constant_(m.weight, 1.0)1159 elif isinstance(m, LayerNorm2d):1160 nn.init.constant_(m.bias, 0)1161 nn.init.constant_(m.weight, 1.0)1162 elif isinstance(m, nn.BatchNorm2d):1163 nn.init.ones_(m.weight)1164 nn.init.zeros_(m.bias)1165 1166 @torch.jit.ignore1167 def no_weight_decay_keywords(self):1168 return {'rpb'}1169 1170 def forward_features(self, x):1171 _, _, H, W = x.shape1172 if H % 32 != 0 or W % 32 != 0:1173 raise ValueError(f"E-RADIO requires input dimensions to be divisible by 32 but got H x W: {H} x {W}")1174 x = self.patch_embed(x)1175 full_features = None1176 for il, level in enumerate(self.levels):1177 x, pre_downsample_x = level(x)1178 1179 if self.return_full_features or self.use_neck:1180 full_features = self.high_res_neck(pre_downsample_x, il, full_features)1181 1182 # x = self.norm(full_features if (self.return_full_features or self.use_neck) else x)1183 x = self.norm(x) # new version for1184 1185 if not self.return_full_features:1186 return x, None1187 1188 return x, full_features1189 1190 def forward(self, x):1191 x, full_features = self.forward_features(x)1192 1193 x = self.avgpool(x)1194 x = torch.flatten(x, 1)1195 1196 x = self.head(x)1197 if full_features is not None:1198 return x, full_features1199 return x1200