wolf6/GARBAGE
0
1import math2from copy import copy3from pathlib import Path4 5import numpy as np6import pandas as pd7import requests8import torch9import torch.nn as nn10import torch.nn.functional as F11from torchvision.ops import DeformConv2d12from PIL import Image13from torch.cuda import amp14 15from utils.datasets import letterbox16from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh17from utils.plots import color_list, plot_one_box18from utils.torch_utils import time_synchronized19 20 21##### basic ####22 23def autopad(k, p=None): # kernel, padding24 # Pad to 'same'25 if p is None:26 p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad27 return p28 29 30class MP(nn.Module):31 def __init__(self, k=2):32 super(MP, self).__init__()33 self.m = nn.MaxPool2d(kernel_size=k, stride=k)34 35 def forward(self, x):36 return self.m(x)37 38 39class SP(nn.Module):40 def __init__(self, k=3, s=1):41 super(SP, self).__init__()42 self.m = nn.MaxPool2d(kernel_size=k, stride=s, padding=k // 2)43 44 def forward(self, x):45 return self.m(x)46 47 48class ReOrg(nn.Module):49 def __init__(self):50 super(ReOrg, self).__init__()51 52 def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)53 return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)54 55 56class Concat(nn.Module):57 def __init__(self, dimension=1):58 super(Concat, self).__init__()59 self.d = dimension60 61 def forward(self, x):62 return torch.cat(x, self.d)63 64 65class Chuncat(nn.Module):66 def __init__(self, dimension=1):67 super(Chuncat, self).__init__()68 self.d = dimension69 70 def forward(self, x):71 x1 = []72 x2 = []73 for xi in x:74 xi1, xi2 = xi.chunk(2, self.d)75 x1.append(xi1)76 x2.append(xi2)77 return torch.cat(x1+x2, self.d)78 79 80class Shortcut(nn.Module):81 def __init__(self, dimension=0):82 super(Shortcut, self).__init__()83 self.d = dimension84 85 def forward(self, x):86 return x[0]+x[1]87 88 89class Foldcut(nn.Module):90 def __init__(self, dimension=0):91 super(Foldcut, self).__init__()92 self.d = dimension93 94 def forward(self, x):95 x1, x2 = x.chunk(2, self.d)96 return x1+x297 98 99class Conv(nn.Module):100 # Standard convolution101 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups102 super(Conv, self).__init__()103 self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)104 self.bn = nn.BatchNorm2d(c2)105 self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())106 107 def forward(self, x):108 return self.act(self.bn(self.conv(x)))109 110 def fuseforward(self, x):111 return self.act(self.conv(x))112 113 114class RobustConv(nn.Module):115 # Robust convolution (use high kernel size 7-11 for: downsampling and other layers). Train for 300 - 450 epochs.116 def __init__(self, c1, c2, k=7, s=1, p=None, g=1, act=True, layer_scale_init_value=1e-6): # ch_in, ch_out, kernel, stride, padding, groups117 super(RobustConv, self).__init__()118 self.conv_dw = Conv(c1, c1, k=k, s=s, p=p, g=c1, act=act)119 self.conv1x1 = nn.Conv2d(c1, c2, 1, 1, 0, groups=1, bias=True)120 self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(c2)) if layer_scale_init_value > 0 else None121 122 def forward(self, x):123 x = x.to(memory_format=torch.channels_last)124 x = self.conv1x1(self.conv_dw(x))125 if self.gamma is not None:126 x = x.mul(self.gamma.reshape(1, -1, 1, 1)) 127 return x128 129 130class RobustConv2(nn.Module):131 # Robust convolution 2 (use [32, 5, 2] or [32, 7, 4] or [32, 11, 8] for one of the paths in CSP).132 def __init__(self, c1, c2, k=7, s=4, p=None, g=1, act=True, layer_scale_init_value=1e-6): # ch_in, ch_out, kernel, stride, padding, groups133 super(RobustConv2, self).__init__()134 self.conv_strided = Conv(c1, c1, k=k, s=s, p=p, g=c1, act=act)135 self.conv_deconv = nn.ConvTranspose2d(in_channels=c1, out_channels=c2, kernel_size=s, stride=s, 136 padding=0, bias=True, dilation=1, groups=1137 )138 self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(c2)) if layer_scale_init_value > 0 else None139 140 def forward(self, x):141 x = self.conv_deconv(self.conv_strided(x))142 if self.gamma is not None:143 x = x.mul(self.gamma.reshape(1, -1, 1, 1)) 144 return x145 146 147def DWConv(c1, c2, k=1, s=1, act=True):148 # Depthwise convolution149 return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)150 151 152class GhostConv(nn.Module):153 # Ghost Convolution https://github.com/huawei-noah/ghostnet154 def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups155 super(GhostConv, self).__init__()156 c_ = c2 // 2 # hidden channels157 self.cv1 = Conv(c1, c_, k, s, None, g, act)158 self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)159 160 def forward(self, x):161 y = self.cv1(x)162 return torch.cat([y, self.cv2(y)], 1)163 164 165class Stem(nn.Module):166 # Stem167 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups168 super(Stem, self).__init__()169 c_ = int(c2/2) # hidden channels170 self.cv1 = Conv(c1, c_, 3, 2)171 self.cv2 = Conv(c_, c_, 1, 1)172 self.cv3 = Conv(c_, c_, 3, 2)173 self.pool = torch.nn.MaxPool2d(2, stride=2)174 self.cv4 = Conv(2 * c_, c2, 1, 1)175 176 def forward(self, x):177 x = self.cv1(x)178 return self.cv4(torch.cat((self.cv3(self.cv2(x)), self.pool(x)), dim=1))179 180 181class DownC(nn.Module):182 # Spatial pyramid pooling layer used in YOLOv3-SPP183 def __init__(self, c1, c2, n=1, k=2):184 super(DownC, self).__init__()185 c_ = int(c1) # hidden channels186 self.cv1 = Conv(c1, c_, 1, 1)187 self.cv2 = Conv(c_, c2//2, 3, k)188 self.cv3 = Conv(c1, c2//2, 1, 1)189 self.mp = nn.MaxPool2d(kernel_size=k, stride=k)190 191 def forward(self, x):192 return torch.cat((self.cv2(self.cv1(x)), self.cv3(self.mp(x))), dim=1)193 194 195class SPP(nn.Module):196 # Spatial pyramid pooling layer used in YOLOv3-SPP197 def __init__(self, c1, c2, k=(5, 9, 13)):198 super(SPP, self).__init__()199 c_ = c1 // 2 # hidden channels200 self.cv1 = Conv(c1, c_, 1, 1)201 self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)202 self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])203 204 def forward(self, x):205 x = self.cv1(x)206 return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))207 208 209class Bottleneck(nn.Module):210 # Darknet bottleneck211 def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion212 super(Bottleneck, self).__init__()213 c_ = int(c2 * e) # hidden channels214 self.cv1 = Conv(c1, c_, 1, 1)215 self.cv2 = Conv(c_, c2, 3, 1, g=g)216 self.add = shortcut and c1 == c2217 218 def forward(self, x):219 return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))220 221 222class Res(nn.Module):223 # ResNet bottleneck224 def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion225 super(Res, self).__init__()226 c_ = int(c2 * e) # hidden channels227 self.cv1 = Conv(c1, c_, 1, 1)228 self.cv2 = Conv(c_, c_, 3, 1, g=g)229 self.cv3 = Conv(c_, c2, 1, 1)230 self.add = shortcut and c1 == c2231 232 def forward(self, x):233 return x + self.cv3(self.cv2(self.cv1(x))) if self.add else self.cv3(self.cv2(self.cv1(x)))234 235 236class ResX(Res):237 # ResNet bottleneck238 def __init__(self, c1, c2, shortcut=True, g=32, e=0.5): # ch_in, ch_out, shortcut, groups, expansion239 super().__init__(c1, c2, shortcut, g, e)240 c_ = int(c2 * e) # hidden channels241 242 243class Ghost(nn.Module):244 # Ghost Bottleneck https://github.com/huawei-noah/ghostnet245 def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride246 super(Ghost, self).__init__()247 c_ = c2 // 2248 self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw249 DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw250 GhostConv(c_, c2, 1, 1, act=False)) # pw-linear251 self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),252 Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()253 254 def forward(self, x):255 return self.conv(x) + self.shortcut(x)256 257##### end of basic #####258 259 260##### cspnet #####261 262class SPPCSPC(nn.Module):263 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks264 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)):265 super(SPPCSPC, self).__init__()266 c_ = int(2 * c2 * e) # hidden channels267 self.cv1 = Conv(c1, c_, 1, 1)268 self.cv2 = Conv(c1, c_, 1, 1)269 self.cv3 = Conv(c_, c_, 3, 1)270 self.cv4 = Conv(c_, c_, 1, 1)271 self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])272 self.cv5 = Conv(4 * c_, c_, 1, 1)273 self.cv6 = Conv(c_, c_, 3, 1)274 self.cv7 = Conv(2 * c_, c2, 1, 1)275 276 def forward(self, x):277 x1 = self.cv4(self.cv3(self.cv1(x)))278 y1 = self.cv6(self.cv5(torch.cat([x1] + [m(x1) for m in self.m], 1)))279 y2 = self.cv2(x)280 return self.cv7(torch.cat((y1, y2), dim=1))281 282class GhostSPPCSPC(SPPCSPC):283 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks284 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)):285 super().__init__(c1, c2, n, shortcut, g, e, k)286 c_ = int(2 * c2 * e) # hidden channels287 self.cv1 = GhostConv(c1, c_, 1, 1)288 self.cv2 = GhostConv(c1, c_, 1, 1)289 self.cv3 = GhostConv(c_, c_, 3, 1)290 self.cv4 = GhostConv(c_, c_, 1, 1)291 self.cv5 = GhostConv(4 * c_, c_, 1, 1)292 self.cv6 = GhostConv(c_, c_, 3, 1)293 self.cv7 = GhostConv(2 * c_, c2, 1, 1)294 295 296class GhostStem(Stem):297 # Stem298 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups299 super().__init__(c1, c2, k, s, p, g, act)300 c_ = int(c2/2) # hidden channels301 self.cv1 = GhostConv(c1, c_, 3, 2)302 self.cv2 = GhostConv(c_, c_, 1, 1)303 self.cv3 = GhostConv(c_, c_, 3, 2)304 self.cv4 = GhostConv(2 * c_, c2, 1, 1)305 306 307class BottleneckCSPA(nn.Module):308 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks309 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion310 super(BottleneckCSPA, self).__init__()311 c_ = int(c2 * e) # hidden channels312 self.cv1 = Conv(c1, c_, 1, 1)313 self.cv2 = Conv(c1, c_, 1, 1)314 self.cv3 = Conv(2 * c_, c2, 1, 1)315 self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])316 317 def forward(self, x):318 y1 = self.m(self.cv1(x))319 y2 = self.cv2(x)320 return self.cv3(torch.cat((y1, y2), dim=1))321 322 323class BottleneckCSPB(nn.Module):324 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks325 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion326 super(BottleneckCSPB, self).__init__()327 c_ = int(c2) # hidden channels328 self.cv1 = Conv(c1, c_, 1, 1)329 self.cv2 = Conv(c_, c_, 1, 1)330 self.cv3 = Conv(2 * c_, c2, 1, 1)331 self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])332 333 def forward(self, x):334 x1 = self.cv1(x)335 y1 = self.m(x1)336 y2 = self.cv2(x1)337 return self.cv3(torch.cat((y1, y2), dim=1))338 339 340class BottleneckCSPC(nn.Module):341 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks342 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion343 super(BottleneckCSPC, self).__init__()344 c_ = int(c2 * e) # hidden channels345 self.cv1 = Conv(c1, c_, 1, 1)346 self.cv2 = Conv(c1, c_, 1, 1)347 self.cv3 = Conv(c_, c_, 1, 1)348 self.cv4 = Conv(2 * c_, c2, 1, 1)349 self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])350 351 def forward(self, x):352 y1 = self.cv3(self.m(self.cv1(x)))353 y2 = self.cv2(x)354 return self.cv4(torch.cat((y1, y2), dim=1))355 356 357class ResCSPA(BottleneckCSPA):358 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks359 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion360 super().__init__(c1, c2, n, shortcut, g, e)361 c_ = int(c2 * e) # hidden channels362 self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=0.5) for _ in range(n)])363 364 365class ResCSPB(BottleneckCSPB):366 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks367 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion368 super().__init__(c1, c2, n, shortcut, g, e)369 c_ = int(c2) # hidden channels370 self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=0.5) for _ in range(n)])371 372 373class ResCSPC(BottleneckCSPC):374 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks375 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion376 super().__init__(c1, c2, n, shortcut, g, e)377 c_ = int(c2 * e) # hidden channels378 self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=0.5) for _ in range(n)])379 380 381class ResXCSPA(ResCSPA):382 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks383 def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion384 super().__init__(c1, c2, n, shortcut, g, e)385 c_ = int(c2 * e) # hidden channels386 self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=1.0) for _ in range(n)])387 388 389class ResXCSPB(ResCSPB):390 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks391 def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion392 super().__init__(c1, c2, n, shortcut, g, e)393 c_ = int(c2) # hidden channels394 self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=1.0) for _ in range(n)])395 396 397class ResXCSPC(ResCSPC):398 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks399 def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion400 super().__init__(c1, c2, n, shortcut, g, e)401 c_ = int(c2 * e) # hidden channels402 self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=1.0) for _ in range(n)])403 404 405class GhostCSPA(BottleneckCSPA):406 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks407 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion408 super().__init__(c1, c2, n, shortcut, g, e)409 c_ = int(c2 * e) # hidden channels410 self.m = nn.Sequential(*[Ghost(c_, c_) for _ in range(n)])411 412 413class GhostCSPB(BottleneckCSPB):414 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks415 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion416 super().__init__(c1, c2, n, shortcut, g, e)417 c_ = int(c2) # hidden channels418 self.m = nn.Sequential(*[Ghost(c_, c_) for _ in range(n)])419 420 421class GhostCSPC(BottleneckCSPC):422 # CSP https://github.com/WongKinYiu/CrossStagePartialNetworks423 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion424 super().__init__(c1, c2, n, shortcut, g, e)425 c_ = int(c2 * e) # hidden channels426 self.m = nn.Sequential(*[Ghost(c_, c_) for _ in range(n)])427 428##### end of cspnet #####429 430 431##### yolor #####432 433class ImplicitA(nn.Module):434 def __init__(self, channel, mean=0., std=.02):435 super(ImplicitA, self).__init__()436 self.channel = channel437 self.mean = mean438 self.std = std439 self.implicit = nn.Parameter(torch.zeros(1, channel, 1, 1))440 nn.init.normal_(self.implicit, mean=self.mean, std=self.std)441 442 def forward(self, x):443 return self.implicit + x444 445 446class ImplicitM(nn.Module):447 def __init__(self, channel, mean=1., std=.02):448 super(ImplicitM, self).__init__()449 self.channel = channel450 self.mean = mean451 self.std = std452 self.implicit = nn.Parameter(torch.ones(1, channel, 1, 1))453 nn.init.normal_(self.implicit, mean=self.mean, std=self.std)454 455 def forward(self, x):456 return self.implicit * x457 458##### end of yolor #####459 460 461##### repvgg #####462 463class RepConv(nn.Module):464 # Represented convolution465 # https://arxiv.org/abs/2101.03697466 467 def __init__(self, c1, c2, k=3, s=1, p=None, g=1, act=True, deploy=False):468 super(RepConv, self).__init__()469 470 self.deploy = deploy471 self.groups = g472 self.in_channels = c1473 self.out_channels = c2474 475 assert k == 3476 assert autopad(k, p) == 1477 478 padding_11 = autopad(k, p) - k // 2479 480 self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())481 482 if deploy:483 self.rbr_reparam = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=True)484 485 else:486 self.rbr_identity = (nn.BatchNorm2d(num_features=c1) if c2 == c1 and s == 1 else None)487 488 self.rbr_dense = nn.Sequential(489 nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False),490 nn.BatchNorm2d(num_features=c2),491 )492 493 self.rbr_1x1 = nn.Sequential(494 nn.Conv2d( c1, c2, 1, s, padding_11, groups=g, bias=False),495 nn.BatchNorm2d(num_features=c2),496 )497 498 def forward(self, inputs):499 if hasattr(self, "rbr_reparam"):500 return self.act(self.rbr_reparam(inputs))501 502 if self.rbr_identity is None:503 id_out = 0504 else:505 id_out = self.rbr_identity(inputs)506 507 return self.act(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out)508 509 def get_equivalent_kernel_bias(self):510 kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)511 kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)512 kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)513 return (514 kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid,515 bias3x3 + bias1x1 + biasid,516 )517 518 def _pad_1x1_to_3x3_tensor(self, kernel1x1):519 if kernel1x1 is None:520 return 0521 else:522 return nn.functional.pad(kernel1x1, [1, 1, 1, 1])523 524 def _fuse_bn_tensor(self, branch):525 if branch is None:526 return 0, 0527 if isinstance(branch, nn.Sequential):528 kernel = branch[0].weight529 running_mean = branch[1].running_mean530 running_var = branch[1].running_var531 gamma = branch[1].weight532 beta = branch[1].bias533 eps = branch[1].eps534 else:535 assert isinstance(branch, nn.BatchNorm2d)536 if not hasattr(self, "id_tensor"):537 input_dim = self.in_channels // self.groups538 kernel_value = np.zeros(539 (self.in_channels, input_dim, 3, 3), dtype=np.float32540 )541 for i in range(self.in_channels):542 kernel_value[i, i % input_dim, 1, 1] = 1543 self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)544 kernel = self.id_tensor545 running_mean = branch.running_mean546 running_var = branch.running_var547 gamma = branch.weight548 beta = branch.bias549 eps = branch.eps550 std = (running_var + eps).sqrt()551 t = (gamma / std).reshape(-1, 1, 1, 1)552 return kernel * t, beta - running_mean * gamma / std553 554 def repvgg_convert(self):555 kernel, bias = self.get_equivalent_kernel_bias()556 return (557 kernel.detach().cpu().numpy(),558 bias.detach().cpu().numpy(),559 )560 561 def fuse_conv_bn(self, conv, bn):562 563 std = (bn.running_var + bn.eps).sqrt()564 bias = bn.bias - bn.running_mean * bn.weight / std565 566 t = (bn.weight / std).reshape(-1, 1, 1, 1)567 weights = conv.weight * t568 569 bn = nn.Identity()570 conv = nn.Conv2d(in_channels = conv.in_channels,571 out_channels = conv.out_channels,572 kernel_size = conv.kernel_size,573 stride=conv.stride,574 padding = conv.padding,575 dilation = conv.dilation,576 groups = conv.groups,577 bias = True,578 padding_mode = conv.padding_mode)579 580 conv.weight = torch.nn.Parameter(weights)581 conv.bias = torch.nn.Parameter(bias)582 return conv583 584 def fuse_repvgg_block(self): 585 if self.deploy:586 return587 print(f"RepConv.fuse_repvgg_block")588 589 self.rbr_dense = self.fuse_conv_bn(self.rbr_dense[0], self.rbr_dense[1])590 591 self.rbr_1x1 = self.fuse_conv_bn(self.rbr_1x1[0], self.rbr_1x1[1])592 rbr_1x1_bias = self.rbr_1x1.bias593 weight_1x1_expanded = torch.nn.functional.pad(self.rbr_1x1.weight, [1, 1, 1, 1])594 595 # Fuse self.rbr_identity596 if (isinstance(self.rbr_identity, nn.BatchNorm2d) or isinstance(self.rbr_identity, nn.modules.batchnorm.SyncBatchNorm)):597 # print(f"fuse: rbr_identity == BatchNorm2d or SyncBatchNorm")598 identity_conv_1x1 = nn.Conv2d(599 in_channels=self.in_channels,600 out_channels=self.out_channels,601 kernel_size=1,602 stride=1,603 padding=0,604 groups=self.groups, 605 bias=False)606 identity_conv_1x1.weight.data = identity_conv_1x1.weight.data.to(self.rbr_1x1.weight.data.device)607 identity_conv_1x1.weight.data = identity_conv_1x1.weight.data.squeeze().squeeze()608 # print(f" identity_conv_1x1.weight = {identity_conv_1x1.weight.shape}")609 identity_conv_1x1.weight.data.fill_(0.0)610 identity_conv_1x1.weight.data.fill_diagonal_(1.0)611 identity_conv_1x1.weight.data = identity_conv_1x1.weight.data.unsqueeze(2).unsqueeze(3)612 # print(f" identity_conv_1x1.weight = {identity_conv_1x1.weight.shape}")613 614 identity_conv_1x1 = self.fuse_conv_bn(identity_conv_1x1, self.rbr_identity)615 bias_identity_expanded = identity_conv_1x1.bias616 weight_identity_expanded = torch.nn.functional.pad(identity_conv_1x1.weight, [1, 1, 1, 1]) 617 else:618 # print(f"fuse: rbr_identity != BatchNorm2d, rbr_identity = {self.rbr_identity}")619 bias_identity_expanded = torch.nn.Parameter( torch.zeros_like(rbr_1x1_bias) )620 weight_identity_expanded = torch.nn.Parameter( torch.zeros_like(weight_1x1_expanded) ) 621 622 623 #print(f"self.rbr_1x1.weight = {self.rbr_1x1.weight.shape}, ")624 #print(f"weight_1x1_expanded = {weight_1x1_expanded.shape}, ")625 #print(f"self.rbr_dense.weight = {self.rbr_dense.weight.shape}, ")626 627 self.rbr_dense.weight = torch.nn.Parameter(self.rbr_dense.weight + weight_1x1_expanded + weight_identity_expanded)628 self.rbr_dense.bias = torch.nn.Parameter(self.rbr_dense.bias + rbr_1x1_bias + bias_identity_expanded)629 630 self.rbr_reparam = self.rbr_dense631 self.deploy = True632 633 if self.rbr_identity is not None:634 del self.rbr_identity635 self.rbr_identity = None636 637 if self.rbr_1x1 is not None:638 del self.rbr_1x1639 self.rbr_1x1 = None640 641 if self.rbr_dense is not None:642 del self.rbr_dense643 self.rbr_dense = None644 645 646class RepBottleneck(Bottleneck):647 # Standard bottleneck648 def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion649 super().__init__(c1, c2, shortcut=True, g=1, e=0.5)650 c_ = int(c2 * e) # hidden channels651 self.cv2 = RepConv(c_, c2, 3, 1, g=g)652 653 654class RepBottleneckCSPA(BottleneckCSPA):655 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks656 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion657 super().__init__(c1, c2, n, shortcut, g, e)658 c_ = int(c2 * e) # hidden channels659 self.m = nn.Sequential(*[RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])660 661 662class RepBottleneckCSPB(BottleneckCSPB):663 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks664 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion665 super().__init__(c1, c2, n, shortcut, g, e)666 c_ = int(c2) # hidden channels667 self.m = nn.Sequential(*[RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])668 669 670class RepBottleneckCSPC(BottleneckCSPC):671 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks672 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion673 super().__init__(c1, c2, n, shortcut, g, e)674 c_ = int(c2 * e) # hidden channels675 self.m = nn.Sequential(*[RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])676 677 678class RepRes(Res):679 # Standard bottleneck680 def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion681 super().__init__(c1, c2, shortcut, g, e)682 c_ = int(c2 * e) # hidden channels683 self.cv2 = RepConv(c_, c_, 3, 1, g=g)684 685 686class RepResCSPA(ResCSPA):687 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks688 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion689 super().__init__(c1, c2, n, shortcut, g, e)690 c_ = int(c2 * e) # hidden channels691 self.m = nn.Sequential(*[RepRes(c_, c_, shortcut, g, e=0.5) for _ in range(n)])692 693 694class RepResCSPB(ResCSPB):695 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks696 def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion697 super().__init__(c1, c2, n, shortcut, g, e)698 c_ = int(c2) # hidden channels699 self.m = nn.Sequential(*[RepRes(c_, c_, shortcut, g, e=0.5) for _ in range(n)])700 701 702class RepResCSPC(ResCSPC):703 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks704 def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion705 super().__init__(c1, c2, n, shortcut, g, e)706 c_ = int(c2 * e) # hidden channels707 self.m = nn.Sequential(*[RepRes(c_, c_, shortcut, g, e=0.5) for _ in range(n)])708 709 710class RepResX(ResX):711 # Standard bottleneck712 def __init__(self, c1, c2, shortcut=True, g=32, e=0.5): # ch_in, ch_out, shortcut, groups, expansion713 super().__init__(c1, c2, shortcut, g, e)714 c_ = int(c2 * e) # hidden channels715 self.cv2 = RepConv(c_, c_, 3, 1, g=g)716 717 718class RepResXCSPA(ResXCSPA):719 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks720 def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion721 super().__init__(c1, c2, n, shortcut, g, e)722 c_ = int(c2 * e) # hidden channels723 self.m = nn.Sequential(*[RepResX(c_, c_, shortcut, g, e=0.5) for _ in range(n)])724 725 726class RepResXCSPB(ResXCSPB):727 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks728 def __init__(self, c1, c2, n=1, shortcut=False, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion729 super().__init__(c1, c2, n, shortcut, g, e)730 c_ = int(c2) # hidden channels731 self.m = nn.Sequential(*[RepResX(c_, c_, shortcut, g, e=0.5) for _ in range(n)])732 733 734class RepResXCSPC(ResXCSPC):735 # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks736 def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion737 super().__init__(c1, c2, n, shortcut, g, e)738 c_ = int(c2 * e) # hidden channels739 self.m = nn.Sequential(*[RepResX(c_, c_, shortcut, g, e=0.5) for _ in range(n)])740 741##### end of repvgg #####742 743 744##### transformer #####745 746class TransformerLayer(nn.Module):747 # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)748 def __init__(self, c, num_heads):749 super().__init__()750 self.q = nn.Linear(c, c, bias=False)751 self.k = nn.Linear(c, c, bias=False)752 self.v = nn.Linear(c, c, bias=False)753 self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)754 self.fc1 = nn.Linear(c, c, bias=False)755 self.fc2 = nn.Linear(c, c, bias=False)756 757 def forward(self, x):758 x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x759 x = self.fc2(self.fc1(x)) + x760 return x761 762 763class TransformerBlock(nn.Module):764 # Vision Transformer https://arxiv.org/abs/2010.11929765 def __init__(self, c1, c2, num_heads, num_layers):766 super().__init__()767 self.conv = None768 if c1 != c2:769 self.conv = Conv(c1, c2)770 self.linear = nn.Linear(c2, c2) # learnable position embedding771 self.tr = nn.Sequential(*[TransformerLayer(c2, num_heads) for _ in range(num_layers)])772 self.c2 = c2773 774 def forward(self, x):775 if self.conv is not None:776 x = self.conv(x)777 b, _, w, h = x.shape778 p = x.flatten(2)779 p = p.unsqueeze(0)780 p = p.transpose(0, 3)781 p = p.squeeze(3)782 e = self.linear(p)783 x = p + e784 785 x = self.tr(x)786 x = x.unsqueeze(3)787 x = x.transpose(0, 3)788 x = x.reshape(b, self.c2, w, h)789 return x790 791##### end of transformer #####792 793 794##### yolov5 #####795 796class Focus(nn.Module):797 # Focus wh information into c-space798 def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups799 super(Focus, self).__init__()800 self.conv = Conv(c1 * 4, c2, k, s, p, g, act)801 # self.contract = Contract(gain=2)802 803 def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)804 return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))805 # return self.conv(self.contract(x))806 807 808class SPPF(nn.Module):809 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher810 def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))811 super().__init__()812 c_ = c1 // 2 # hidden channels813 self.cv1 = Conv(c1, c_, 1, 1)814 self.cv2 = Conv(c_ * 4, c2, 1, 1)815 self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)816 817 def forward(self, x):818 x = self.cv1(x)819 y1 = self.m(x)820 y2 = self.m(y1)821 return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))822 823 824class Contract(nn.Module):825 # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)826 def __init__(self, gain=2):827 super().__init__()828 self.gain = gain829 830 def forward(self, x):831 N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'832 s = self.gain833 x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2)834 x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)835 return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40)836 837 838class Expand(nn.Module):839 # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)840 def __init__(self, gain=2):841 super().__init__()842 self.gain = gain843 844 def forward(self, x):845 N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'846 s = self.gain847 x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80)848 x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)849 return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160)850 851 852class NMS(nn.Module):853 # Non-Maximum Suppression (NMS) module854 conf = 0.25 # confidence threshold855 iou = 0.45 # IoU threshold856 classes = None # (optional list) filter by class857 858 def __init__(self):859 super(NMS, self).__init__()860 861 def forward(self, x):862 return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)863 864 865class autoShape(nn.Module):866 # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS867 conf = 0.25 # NMS confidence threshold868 iou = 0.45 # NMS IoU threshold869 classes = None # (optional list) filter by class870 871 def __init__(self, model):872 super(autoShape, self).__init__()873 self.model = model.eval()874 875 def autoshape(self):876 print('autoShape already enabled, skipping... ') # model already converted to model.autoshape()877 return self878 879 @torch.no_grad()880 def forward(self, imgs, size=640, augment=False, profile=False):881 # Inference from various sources. For height=640, width=1280, RGB images example inputs are:882 # filename: imgs = 'data/samples/zidane.jpg'883 # URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'884 # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)885 # PIL: = Image.open('image.jpg') # HWC x(640,1280,3)886 # numpy: = np.zeros((640,1280,3)) # HWC887 # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)888 # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images889 890 t = [time_synchronized()]891 p = next(self.model.parameters()) # for device and type892 if isinstance(imgs, torch.Tensor): # torch893 with amp.autocast(enabled=p.device.type != 'cpu'):894 return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference895 896 # Pre-process897 n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images898 shape0, shape1, files = [], [], [] # image and inference shapes, filenames899 for i, im in enumerate(imgs):900 f = f'image{i}' # filename901 if isinstance(im, str): # filename or uri902 im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im903 elif isinstance(im, Image.Image): # PIL Image904 im, f = np.asarray(im), getattr(im, 'filename', f) or f905 files.append(Path(f).with_suffix('.jpg').name)906 if im.shape[0] < 5: # image in CHW907 im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)908 im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input909 s = im.shape[:2] # HWC910 shape0.append(s) # image shape911 g = (size / max(s)) # gain912 shape1.append([y * g for y in s])913 imgs[i] = im # update914 shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape915 x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad916 x = np.stack(x, 0) if n > 1 else x[0][None] # stack917 x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW918 x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32919 t.append(time_synchronized())920 921 with amp.autocast(enabled=p.device.type != 'cpu'):922 # Inference923 y = self.model(x, augment, profile)[0] # forward924 t.append(time_synchronized())925 926 # Post-process927 y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS928 for i in range(n):929 scale_coords(shape1, y[i][:, :4], shape0[i])930 931 t.append(time_synchronized())932 return Detections(imgs, y, files, t, self.names, x.shape)933 934 935class Detections:936 # detections class for YOLOv5 inference results937 def __init__(self, imgs, pred, files, times=None, names=None, shape=None):938 super(Detections, self).__init__()939 d = pred[0].device # device940 gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations941 self.imgs = imgs # list of images as numpy arrays942 self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)943 self.names = names # class names944 self.files = files # image filenames945 self.xyxy = pred # xyxy pixels946 self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels947 self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized948 self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized949 self.n = len(self.pred) # number of images (batch size)950 self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)951 self.s = shape # inference BCHW shape952 953 def display(self, pprint=False, show=False, save=False, render=False, save_dir=''):954 colors = color_list()955 for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):956 str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '957 if pred is not None:958 for c in pred[:, -1].unique():959 n = (pred[:, -1] == c).sum() # detections per class960 str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string961 if show or save or render:962 for *box, conf, cls in pred: # xyxy, confidence, class963 label = f'{self.names[int(cls)]} {conf:.2f}'964 plot_one_box(box, img, label=label, color=colors[int(cls) % 10])965 img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np966 if pprint:967 print(str.rstrip(', '))968 if show:969 img.show(self.files[i]) # show970 if save:971 f = self.files[i]972 img.save(Path(save_dir) / f) # save973 print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')974 if render:975 self.imgs[i] = np.asarray(img)976 977 def print(self):978 self.display(pprint=True) # print results979 print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t)980 981 def show(self):982 self.display(show=True) # show results983 984 def save(self, save_dir='runs/hub/exp'):985 save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp') # increment save_dir986 Path(save_dir).mkdir(parents=True, exist_ok=True)987 self.display(save=True, save_dir=save_dir) # save results988 989 def render(self):990 self.display(render=True) # render results991 return self.imgs992 993 def pandas(self):994 # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])995 new = copy(self) # return copy996 ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns997 cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns998 for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):999 a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update1000 setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])1001 return new1002 1003 def tolist(self):1004 # return a list of Detections objects, i.e. 'for result in results.tolist():'1005 x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)]1006 for d in x:1007 for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:1008 setattr(d, k, getattr(d, k)[0]) # pop out of list1009 return x1010 1011 def __len__(self):1012 return self.n1013 1014 1015class Classify(nn.Module):1016 # Classification head, i.e. x(b,c1,20,20) to x(b,c2)1017 def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups1018 super(Classify, self).__init__()1019 self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)1020 self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)1021 self.flat = nn.Flatten()1022 1023 def forward(self, x):1024 z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list1025 return self.flat(self.conv(z)) # flatten to x(b,c2)1026 1027##### end of yolov5 ######1028 1029 1030##### orepa #####1031 1032def transI_fusebn(kernel, bn):1033 gamma = bn.weight1034 std = (bn.running_var + bn.eps).sqrt()1035 return kernel * ((gamma / std).reshape(-1, 1, 1, 1)), bn.bias - bn.running_mean * gamma / std1036 1037 1038class ConvBN(nn.Module):1039 def __init__(self, in_channels, out_channels, kernel_size,1040 stride=1, padding=0, dilation=1, groups=1, deploy=False, nonlinear=None):1041 super().__init__()1042 if nonlinear is None:1043 self.nonlinear = nn.Identity()1044 else:1045 self.nonlinear = nonlinear1046 if deploy:1047 self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,1048 stride=stride, padding=padding, dilation=dilation, groups=groups, bias=True)1049 else:1050 self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,1051 stride=stride, padding=padding, dilation=dilation, groups=groups, bias=False)1052 self.bn = nn.BatchNorm2d(num_features=out_channels)1053 1054 def forward(self, x):1055 if hasattr(self, 'bn'):1056 return self.nonlinear(self.bn(self.conv(x)))1057 else:1058 return self.nonlinear(self.conv(x))1059 1060 def switch_to_deploy(self):1061 kernel, bias = transI_fusebn(self.conv.weight, self.bn)1062 conv = nn.Conv2d(in_channels=self.conv.in_channels, out_channels=self.conv.out_channels, kernel_size=self.conv.kernel_size,1063 stride=self.conv.stride, padding=self.conv.padding, dilation=self.conv.dilation, groups=self.conv.groups, bias=True)1064 conv.weight.data = kernel1065 conv.bias.data = bias1066 for para in self.parameters():1067 para.detach_()1068 self.__delattr__('conv')1069 self.__delattr__('bn')1070 self.conv = conv 1071 1072class OREPA_3x3_RepConv(nn.Module):1073 1074 def __init__(self, in_channels, out_channels, kernel_size,1075 stride=1, padding=0, dilation=1, groups=1,1076 internal_channels_1x1_3x3=None,1077 deploy=False, nonlinear=None, single_init=False):1078 super(OREPA_3x3_RepConv, self).__init__()1079 self.deploy = deploy1080 1081 if nonlinear is None:1082 self.nonlinear = nn.Identity()1083 else:1084 self.nonlinear = nonlinear1085 1086 self.kernel_size = kernel_size1087 self.in_channels = in_channels1088 self.out_channels = out_channels1089 self.groups = groups1090 assert padding == kernel_size // 21091 1092 self.stride = stride1093 self.padding = padding1094 self.dilation = dilation1095 1096 self.branch_counter = 01097 1098 self.weight_rbr_origin = nn.Parameter(torch.Tensor(out_channels, int(in_channels/self.groups), kernel_size, kernel_size))1099 nn.init.kaiming_uniform_(self.weight_rbr_origin, a=math.sqrt(1.0))1100 self.branch_counter += 11101 1102 1103 if groups < out_channels:1104 self.weight_rbr_avg_conv = nn.Parameter(torch.Tensor(out_channels, int(in_channels/self.groups), 1, 1))1105 self.weight_rbr_pfir_conv = nn.Parameter(torch.Tensor(out_channels, int(in_channels/self.groups), 1, 1))1106 nn.init.kaiming_uniform_(self.weight_rbr_avg_conv, a=1.0)1107 nn.init.kaiming_uniform_(self.weight_rbr_pfir_conv, a=1.0)1108 self.weight_rbr_avg_conv.data1109 self.weight_rbr_pfir_conv.data1110 self.register_buffer('weight_rbr_avg_avg', torch.ones(kernel_size, kernel_size).mul(1.0/kernel_size/kernel_size))1111 self.branch_counter += 11112 1113 else:1114 raise NotImplementedError1115 self.branch_counter += 11116 1117 if internal_channels_1x1_3x3 is None:1118 internal_channels_1x1_3x3 = in_channels if groups < out_channels else 2 * in_channels # For mobilenet, it is better to have 2X internal channels1119 1120 if internal_channels_1x1_3x3 == in_channels:1121 self.weight_rbr_1x1_kxk_idconv1 = nn.Parameter(torch.zeros(in_channels, int(in_channels/self.groups), 1, 1))1122 id_value = np.zeros((in_channels, int(in_channels/self.groups), 1, 1))1123 for i in range(in_channels):1124 id_value[i, i % int(in_channels/self.groups), 0, 0] = 11125 id_tensor = torch.from_numpy(id_value).type_as(self.weight_rbr_1x1_kxk_idconv1)1126 self.register_buffer('id_tensor', id_tensor)1127 1128 else:1129 self.weight_rbr_1x1_kxk_conv1 = nn.Parameter(torch.Tensor(internal_channels_1x1_3x3, int(in_channels/self.groups), 1, 1))1130 nn.init.kaiming_uniform_(self.weight_rbr_1x1_kxk_conv1, a=math.sqrt(1.0))1131 self.weight_rbr_1x1_kxk_conv2 = nn.Parameter(torch.Tensor(out_channels, int(internal_channels_1x1_3x3/self.groups), kernel_size, kernel_size))1132 nn.init.kaiming_uniform_(self.weight_rbr_1x1_kxk_conv2, a=math.sqrt(1.0))1133 self.branch_counter += 11134 1135 expand_ratio = 81136 self.weight_rbr_gconv_dw = nn.Parameter(torch.Tensor(in_channels*expand_ratio, 1, kernel_size, kernel_size))1137 self.weight_rbr_gconv_pw = nn.Parameter(torch.Tensor(out_channels, in_channels*expand_ratio, 1, 1))1138 nn.init.kaiming_uniform_(self.weight_rbr_gconv_dw, a=math.sqrt(1.0))1139 nn.init.kaiming_uniform_(self.weight_rbr_gconv_pw, a=math.sqrt(1.0))1140 self.branch_counter += 11141 1142 if out_channels == in_channels and stride == 1:1143 self.branch_counter += 11144 1145 self.vector = nn.Parameter(torch.Tensor(self.branch_counter, self.out_channels))1146 self.bn = nn.BatchNorm2d(out_channels)1147 1148 self.fre_init()1149 1150 nn.init.constant_(self.vector[0, :], 0.25) #origin1151 nn.init.constant_(self.vector[1, :], 0.25) #avg1152 nn.init.constant_(self.vector[2, :], 0.0) #prior1153 nn.init.constant_(self.vector[3, :], 0.5) #1x1_kxk1154 nn.init.constant_(self.vector[4, :], 0.5) #dws_conv1155 1156 1157 def fre_init(self):1158 prior_tensor = torch.Tensor(self.out_channels, self.kernel_size, self.kernel_size)1159 half_fg = self.out_channels/21160 for i in range(self.out_channels):1161 for h in range(3):1162 for w in range(3):1163 if i < half_fg:1164 prior_tensor[i, h, w] = math.cos(math.pi*(h+0.5)*(i+1)/3)1165 else:1166 prior_tensor[i, h, w] = math.cos(math.pi*(w+0.5)*(i+1-half_fg)/3)1167 1168 self.register_buffer('weight_rbr_prior', prior_tensor)1169 1170 def weight_gen(self):1171 1172 weight_rbr_origin = torch.einsum('oihw,o->oihw', self.weight_rbr_origin, self.vector[0, :])1173 1174 weight_rbr_avg = torch.einsum('oihw,o->oihw', torch.einsum('oihw,hw->oihw', self.weight_rbr_avg_conv, self.weight_rbr_avg_avg), self.vector[1, :])1175 1176 weight_rbr_pfir = torch.einsum('oihw,o->oihw', torch.einsum('oihw,ohw->oihw', self.weight_rbr_pfir_conv, self.weight_rbr_prior), self.vector[2, :])1177 1178 weight_rbr_1x1_kxk_conv1 = None1179 if hasattr(self, 'weight_rbr_1x1_kxk_idconv1'):1180 weight_rbr_1x1_kxk_conv1 = (self.weight_rbr_1x1_kxk_idconv1 + self.id_tensor).squeeze()1181 elif hasattr(self, 'weight_rbr_1x1_kxk_conv1'):1182 weight_rbr_1x1_kxk_conv1 = self.weight_rbr_1x1_kxk_conv1.squeeze()1183 else:1184 raise NotImplementedError1185 weight_rbr_1x1_kxk_conv2 = self.weight_rbr_1x1_kxk_conv21186 1187 if self.groups > 1:1188 g = self.groups1189 t, ig = weight_rbr_1x1_kxk_conv1.size()1190 o, tg, h, w = weight_rbr_1x1_kxk_conv2.size()1191 weight_rbr_1x1_kxk_conv1 = weight_rbr_1x1_kxk_conv1.view(g, int(t/g), ig)1192 weight_rbr_1x1_kxk_conv2 = weight_rbr_1x1_kxk_conv2.view(g, int(o/g), tg, h, w)1193 weight_rbr_1x1_kxk = torch.einsum('gti,gothw->goihw', weight_rbr_1x1_kxk_conv1, weight_rbr_1x1_kxk_conv2).view(o, ig, h, w)1194 else:1195 weight_rbr_1x1_kxk = torch.einsum('ti,othw->oihw', weight_rbr_1x1_kxk_conv1, weight_rbr_1x1_kxk_conv2)1196 1197 weight_rbr_1x1_kxk = torch.einsum('oihw,o->oihw', weight_rbr_1x1_kxk, self.vector[3, :])1198 1199 weight_rbr_gconv = self.dwsc2full(self.weight_rbr_gconv_dw, self.weight_rbr_gconv_pw, self.in_channels)1200 weight_rbr_gconv = torch.einsum('oihw,o->oihw', weight_rbr_gconv, self.vector[4, :])