import torch import torch.nn as nn import numpy as np from torchsummary import summary import torch.nn.functional as F class convBatchReluBlock(nn.Module): def __init__(self, in_c, out_c, k, s, p): super().__init__() self.cblr = nn.Sequential( nn.Conv2d(in_c, out_c, kernel_size=k, stride=s, padding=p), nn.BatchNorm2d(out_c, momentum=0.03, eps=1E-4), nn.LeakyReLU(0.1, inplace=True) ) def forward(self, x): return self.cblr(x) class yolo3tiny(nn.Module): def __init__(self): super().__init__() self.con1 = convBatchReluBlock(3,16,3,1,1) self.max1 = nn.MaxPool2d(2,2,0) self.con2 = convBatchReluBlock(16,32,3,1,1) self.max2 = nn.MaxPool2d(2,2,0) self.con3 = convBatchReluBlock(32,64,3,1,1) self.max3 = nn.MaxPool2d(2,2,0) self.con4 = convBatchReluBlock(64,128,3,1,1) self.max4 = nn.MaxPool2d(2,2,0) self.con5 = convBatchReluBlock(128,256,3,1,1)#开始分支 #分支1 继续提取深层的语义信息 并在深层的特征上检测相对于大的目标 self.maxA0 = nn.MaxPool2d(2,2,0) self.conA1 = convBatchReluBlock(256,512,3,1,1) self.zeropad = nn.ZeroPad2d((0, 1, 0, 1)) self.maxA1 = nn.MaxPool2d(2,1,0) self.conA2 = convBatchReluBlock(512,1024,3,1,1) self.conA3 = convBatchReluBlock(1024,256,1,1,0) self.conA4 = convBatchReluBlock(256, 512,3,1,1) self.yolo1 = convBatchReluBlock(512,255,1,1,0) # 分支2 将 self.conA3 特征层 使用反卷积进行上采样之后 与 self.con5 层的特征进行融合 在这一层次检测相对较小的物体 self.conB1 = convBatchReluBlock(256,128,1,1,0) self.upsample = nn.Upsample(scale_factor=2) #融合之后 256 + 128 = 384 self.conB2 = convBatchReluBlock(384,256,3,1,1) self.yolo2 = convBatchReluBlock(256,255,1,1,0) def forward(self, x): x = self.con1(x) x = self.max1(x) x = self.max2(self.con2(x)) x = self.max3(self.con3(x)) x = self.max4(self.con4(x)) con5 = self.con5(x) #分支1 y1 = self.maxA0(con5) y1 = self.maxA1(self.zeropad(self.conA1(y1))) con_A3 = self.conA3(self.conA2(y1)) y1 = self.conA4(con_A3) y1 = self.yolo1(y1) #分支2 y2 = self.conB1(con_A3) y2 = self.upsample(y2) y2 = torch.cat((y2,con5),1) y2 = self.conB2(y2) y2 = self.yolo2(y2) return (y1, y2) if __name__ == '__main__': x = torch.randn(1, 3, 416, 416) model = yolo3tiny() y1, y2 = model(x) print(y1.shape, y2.shape) summary(model, (3, 416, 416),device="cpu")
y1和y2的形状
torch.Size([1, 255, 13, 13]) torch.Size([1, 255, 26, 26])
---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [-1, 16, 416, 416] 448 BatchNorm2d-2 [-1, 16, 416, 416] 32 LeakyReLU-3 [-1, 16, 416, 416] 0 convBatchReluBlock-4 [-1, 16, 416, 416] 0 MaxPool2d-5 [-1, 16, 208, 208] 0 Conv2d-6 [-1, 32, 208, 208] 4,640 BatchNorm2d-7 [-1, 32, 208, 208] 64 LeakyReLU-8 [-1, 32, 208, 208] 0 convBatchReluBlock-9 [-1, 32, 208, 208] 0 MaxPool2d-10 [-1, 32, 104, 104] 0 Conv2d-11 [-1, 64, 104, 104] 18,496 BatchNorm2d-12 [-1, 64, 104, 104] 128 LeakyReLU-13 [-1, 64, 104, 104] 0 convBatchReluBlock-14 [-1, 64, 104, 104] 0 MaxPool2d-15 [-1, 64, 52, 52] 0 Conv2d-16 [-1, 128, 52, 52] 73,856 BatchNorm2d-17 [-1, 128, 52, 52] 256 LeakyReLU-18 [-1, 128, 52, 52] 0 convBatchReluBlock-19 [-1, 128, 52, 52] 0 MaxPool2d-20 [-1, 128, 26, 26] 0 Conv2d-21 [-1, 256, 26, 26] 295,168 BatchNorm2d-22 [-1, 256, 26, 26] 512 LeakyReLU-23 [-1, 256, 26, 26] 0 convBatchReluBlock-24 [-1, 256, 26, 26] 0 MaxPool2d-25 [-1, 256, 13, 13] 0 Conv2d-26 [-1, 512, 13, 13] 1,180,160 BatchNorm2d-27 [-1, 512, 13, 13] 1,024 LeakyReLU-28 [-1, 512, 13, 13] 0 convBatchReluBlock-29 [-1, 512, 13, 13] 0 ZeroPad2d-30 [-1, 512, 14, 14] 0 MaxPool2d-31 [-1, 512, 13, 13] 0 Conv2d-32 [-1, 1024, 13, 13] 4,719,616 BatchNorm2d-33 [-1, 1024, 13, 13] 2,048 LeakyReLU-34 [-1, 1024, 13, 13] 0 convBatchReluBlock-35 [-1, 1024, 13, 13] 0 Conv2d-36 [-1, 256, 13, 13] 262,400 BatchNorm2d-37 [-1, 256, 13, 13] 512 LeakyReLU-38 [-1, 256, 13, 13] 0 convBatchReluBlock-39 [-1, 256, 13, 13] 0 Conv2d-40 [-1, 512, 13, 13] 1,180,160 BatchNorm2d-41 [-1, 512, 13, 13] 1,024 LeakyReLU-42 [-1, 512, 13, 13] 0 convBatchReluBlock-43 [-1, 512, 13, 13] 0 Conv2d-44 [-1, 255, 13, 13] 130,815 BatchNorm2d-45 [-1, 255, 13, 13] 510 LeakyReLU-46 [-1, 255, 13, 13] 0 convBatchReluBlock-47 [-1, 255, 13, 13] 0 Conv2d-48 [-1, 128, 13, 13] 32,896 BatchNorm2d-49 [-1, 128, 13, 13] 256 LeakyReLU-50 [-1, 128, 13, 13] 0 convBatchReluBlock-51 [-1, 128, 13, 13] 0 Upsample-52 [-1, 128, 26, 26] 0 Conv2d-53 [-1, 256, 26, 26] 884,992 BatchNorm2d-54 [-1, 256, 26, 26] 512 LeakyReLU-55 [-1, 256, 26, 26] 0 convBatchReluBlock-56 [-1, 256, 26, 26] 0 Conv2d-57 [-1, 255, 26, 26] 65,535 BatchNorm2d-58 [-1, 255, 26, 26] 510 LeakyReLU-59 [-1, 255, 26, 26] 0 convBatchReluBlock-60 [-1, 255, 26, 26] 0 ================================================================ Total params: 8,856,570 Trainable params: 8,856,570 Non-trainable params: 0 ---------------------------------------------------------------- Input size (MB): 1.98 Forward/backward pass size (MB): 200.44 Params size (MB): 33.79 Estimated Total Size (MB): 236.20 ----------------------------------------------------------------
