python|Yolo4-Tiny代码解读 pytorch|深度学习|神经网络|算法

本文参考代码链接代码
yolo4-tiny的整体结构，如下图所示

文章图片

从图中可以看到，yolo4 tiny大致可以分为以下几个模块：
1.CSPDarknet53-tiny，是整个网络的backbone
2.FPN
3.yolo head
下面分别作一介绍：
CSPDarknet53 tiny 由3个BasicConv块和3个Resblock_body块构成。
BasicConv比较简单，其实就是一个卷积+BN+LeakyReLU的结构，代码如下所示：

class BasicConv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride=1): super(BasicConv, self).__init__()self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, kernel_size//2, bias=False) self.bn = nn.BatchNorm2d(out_channels) self.activation = nn.LeakyReLU(0.1)def forward(self, x): x = self.conv(x) x = self.bn(x) x = self.activation(x) return x

Resblock_body较为复杂，如下图所示

文章图片

它由4个BasicConv块、1个MaxPooling块构成，输入数据经过conv1后，分为两部分，后半部分经过conv2计算后，得到route1，route1进一步经过conv3，计算得到结果，然后再将该结果和route1进行拼接(concat)，输入conv4，计算得到feat，然后将route和feat进行拼接，得到的结果输入maxpool后得到最终的输出。
以下是具体的代码：

#---------------------------------------------------# #CSPdarknet53-tiny的结构块 #存在一个大残差边 #这个大残差边绕过了很多的残差结构 #---------------------------------------------------# class Resblock_body(nn.Module): def __init__(self, in_channels, out_channels): super(Resblock_body, self).__init__() self.out_channels = out_channelsself.conv1 = BasicConv(in_channels, out_channels, 3)self.conv2 = BasicConv(out_channels//2, out_channels//2, 3) self.conv3 = BasicConv(out_channels//2, out_channels//2, 3)self.conv4 = BasicConv(out_channels, out_channels, 1) self.maxpool = nn.MaxPool2d([2,2],[2,2])def forward(self, x): # 利用一个3x3卷积进行特征整合 x = self.conv1(x) # 引出一个大的残差边route route = xc = self.out_channels # 对特征层的通道进行分割，取第二部分作为主干部分。 x = torch.split(x, c//2, dim = 1)[1] # 对主干部分进行3x3卷积 x = self.conv2(x) # 引出一个小的残差边route_1 route1 = x # 对第主干部分进行3x3卷积 x = self.conv3(x) # 主干部分与残差部分进行相接 x = torch.cat([x,route1], dim = 1) # 对相接后的结果进行1x1卷积 x = self.conv4(x) feat = x x = torch.cat([route, x], dim = 1)# 利用最大池化进行高和宽的压缩 x = self.maxpool(x) return x,feat

可以看到，resblock_body块有两个输出，一个是feat，另一个是池化后得到的最终结果。
最终，根据Yolo4 tiny总体的结构图，我们可以得到backbone的具体代码：

class CSPDarkNet(nn.Module): def __init__(self): super(CSPDarkNet, self).__init__() # 首先利用两次步长为2x2的3x3卷积进行高和宽的压缩 # 416,416,3 -> 208,208,32 -> 104,104,64 self.conv1 = BasicConv(3, 32, kernel_size=3, stride=2) self.conv2 = BasicConv(32, 64, kernel_size=3, stride=2)# 104,104,64 -> 52,52,128 self.resblock_body1 =Resblock_body(64, 64) # 52,52,128 -> 26,26,256 self.resblock_body2 =Resblock_body(128, 128) # 26,26,256 -> 13,13,512 self.resblock_body3 =Resblock_body(256, 256) # 13,13,512 -> 13,13,512 self.conv3 = BasicConv(512, 512, kernel_size=3)self.num_features = 1 # 进行权值初始化 for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_()def forward(self, x): # 416,416,3 -> 208,208,32 -> 104,104,64 x = self.conv1(x) x = self.conv2(x)# 104,104,64 -> 52,52,128 x, _= self.resblock_body1(x) # 52,52,128 -> 26,26,256 x, _= self.resblock_body2(x) # 26,26,256 -> x为13,13,512 #-> feat1为26,26,256 x, feat1= self.resblock_body3(x)# 13,13,512 -> 13,13,512 x = self.conv3(x) feat2 = x return feat1,feat2

该backbone有两个输出，分别为26x26x256的feat1和13x13x512的feat2
FPN FPN的详细结构，如下图所示

文章图片
它包括conv_for_P5、upsample和concat等操作，feat2经过conv_for_P5后，得到256x13x13的输出P5，P5一方面输入至yolo_headP5进一步处理，另一方面经过upsample操作（包括一个卷积操作和一个上采样操作）和feat1拼接，得到384x26x26的P4。
以下是upsample模块的具体代码：

class Upsample(nn.Module): def __init__(self, in_channels, out_channels): super(Upsample, self).__init__()self.upsample = nn.Sequential( BasicConv(in_channels, out_channels, 1), nn.Upsample(scale_factor=2, mode='nearest') )def forward(self, x,): x = self.upsample(x) return x

Yolo Head yolo head由两个卷积层、一个激活函数层构成，如下图所示

文章图片

代码：

def yolo_head(filters_list, in_filters): m = nn.Sequential( BasicConv(in_filters, filters_list[0], 3), nn.Conv2d(filters_list[0], filters_list[1], 1), ) return m

Yolo4 Tiny 经过上述分析，不难得到总的yolo4 tiny代码

#---------------------------------------------------# #yolo_body #---------------------------------------------------# class YoloBody(nn.Module): def __init__(self, anchors_mask, num_classes, phi=0, pretrained=False): super(YoloBody, self).__init__() self.phi= phi self.backbone= darknet53_tiny(pretrained)self.conv_for_P5= BasicConv(512,256,1) self.yolo_headP5= yolo_head([512, len(anchors_mask[0]) * (5 + num_classes)],256)self.upsample= Upsample(256,128) self.yolo_headP4= yolo_head([256, len(anchors_mask[1]) * (5 + num_classes)],384)def forward(self, x): #---------------------------------------------------# #生成CSPdarknet53_tiny的主干模型 #feat1的shape为26,26,256 #feat2的shape为13,13,512 #---------------------------------------------------# feat1, feat2 = self.backbone(x)# 13,13,512 -> 13,13,256 P5 = self.conv_for_P5(feat2) # 13,13,256 -> 13,13,512 -> 13,13,255 out0 = self.yolo_headP5(P5) # 13,13,256 -> 13,13,128 -> 26,26,128 P5_Upsample = self.upsample(P5) # 26,26,256 + 26,26,128 -> 26,26,384P4 = torch.cat([P5_Upsample,feat1],axis=1)# 26,26,384 -> 26,26,256 -> 26,26,255 out1 = self.yolo_headP4(P4)return out0, out1

为了加深理解，博主还手动实现了一个yolo4 tiny前向推理的代码

import torch import numpy as np import torch.nn.functional as Fdef FusedBasicConv(x,w,b,s,p=1): x=F.conv2d(input=x,weight=w,bias=b,stride=s,padding=p) return F.leaky_relu(input=x,negative_slope=0.1)def FusedResBlock(x,n,c,w1,w2,w3,w4,b1,b2,b3,b4): x=F.conv2d(input=x,weight=w1,bias=b1,stride=1,padding=1) x=F.leaky_relu(x,negative_slope=0.1) route=x # x=F.conv2d(input=x[:,n//2:n,:,:],weight=w2,bias=b2,stride=1,padding=1) route1=F.leaky_relu(input=x,negative_slope=0.1) # x=F.conv2d(input=route1,weight=w3,bias=b3,stride=1,padding=1) x=F.leaky_relu(input=x,negative_slope=0.1) # x=torch.cat([x,route1],dim=1) # x=F.conv2d(input=x,weight=w4,bias=b4,stride=1,padding=0) feat=F.leaky_relu(input=x,negative_slope=0.1) # x=torch.cat([route,feat],dim=1) x=F.max_pool2d(x,kernel_size=2,stride=2) return x,featclass MyCSPdarknet53_tiny: def __init__(self): self.basic_conv1_w=None self.basic_conv1_b=None self.basic_conv2_w=None self.basic_conv2_b=None self.basic_conv3_w=None self.basic_conv3_b=Nonedef load_weight(self,dir): #basic self.basic_conv1_w = torch.from_numpy(np.fromfile(dir + "\\BasicConv1\\w.bin",dtype=np.float32)).view(32,3,3,3) self.basic_conv1_b = torch.from_numpy(np.fromfile(dir + "\\BasicConv1\\b.bin",dtype=np.float32)) self.basic_conv2_w = torch.from_numpy(np.fromfile(dir + "\\BasicConv2\\w.bin", dtype=np.float32)).view(64,32,3, 3) self.basic_conv2_b = torch.from_numpy(np.fromfile(dir + "\\BasicConv2\\b.bin", dtype=np.float32)) self.basic_conv3_w = torch.from_numpy(np.fromfile(dir + "\\BasicConv3\\w.bin", dtype=np.float32)).view(512,512,3, 3) self.basic_conv3_b = torch.from_numpy(np.fromfile(dir + "\\BasicConv3\\b.bin", dtype=np.float32)) #resblock1,64,64 self.resblock1_w1 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\w1.bin",dtype=np.float32)).view(64,64,3,3) self.resblock1_b1 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\b1.bin",dtype=np.float32)) self.resblock1_w2 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\w2.bin", dtype=np.float32)).view(32,32,3,3) self.resblock1_b2 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\b2.bin", dtype=np.float32)) self.resblock1_w3 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\w3.bin", dtype=np.float32)).view(32,32,3,3) self.resblock1_b3 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\b3.bin", dtype=np.float32)) self.resblock1_w4 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\w4.bin", dtype=np.float32)).view(64,64,1,1) self.resblock1_b4 = torch.from_numpy(np.fromfile(dir + "\\ResBlock1\\b4.bin", dtype=np.float32)) #resblock2,128,128 self.resblock2_w1 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\w1.bin", dtype=np.float32)).view(128,128,3,3) self.resblock2_b1 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\b1.bin", dtype=np.float32)) self.resblock2_w2 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\w2.bin", dtype=np.float32)).view(64,64,3,3) self.resblock2_b2 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\b2.bin", dtype=np.float32)) self.resblock2_w3 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\w3.bin", dtype=np.float32)).view(64,64,3,3) self.resblock2_b3 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\b3.bin", dtype=np.float32)) self.resblock2_w4 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\w4.bin", dtype=np.float32)).view(128,128,1,1) self.resblock2_b4 = torch.from_numpy(np.fromfile(dir + "\\ResBlock2\\b4.bin", dtype=np.float32)) #resblock3,256,256 self.resblock3_w1 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\w1.bin", dtype=np.float32)).view(256,256,3,3) self.resblock3_b1 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\b1.bin", dtype=np.float32)) self.resblock3_w2 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\w2.bin", dtype=np.float32)).view(128,128,3,3) self.resblock3_b2 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\b2.bin", dtype=np.float32)) self.resblock3_w3 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\w3.bin", dtype=np.float32)).view(128,128,3,3) self.resblock3_b3 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\b3.bin", dtype=np.float32)) self.resblock3_w4 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\w4.bin", dtype=np.float32)).view(256,256,1,1) self.resblock3_b4 = torch.from_numpy(np.fromfile(dir + "\\ResBlock3\\b4.bin", dtype=np.float32))def forward(self,x): x=FusedBasicConv(x,self.basic_conv1_w,self.basic_conv1_b,2) x=FusedBasicConv(x,self.basic_conv2_w,self.basic_conv2_b,2) x,_=FusedResBlock(x,64,64, self.resblock1_w1,self.resblock1_w2,self.resblock1_w3,self.resblock1_w4, self.resblock1_b1,self.resblock1_b2,self.resblock1_b3,self.resblock1_b4) x,_=FusedResBlock(x,128,128, self.resblock2_w1, self.resblock2_w2, self.resblock2_w3, self.resblock2_w4, self.resblock2_b1, self.resblock2_b2, self.resblock2_b3, self.resblock2_b4) x,feat1=FusedResBlock(x,256,256, self.resblock3_w1, self.resblock3_w2, self.resblock3_w3, self.resblock3_w4, self.resblock3_b1, self.resblock3_b2, self.resblock3_b3, self.resblock3_b4) feat2=FusedBasicConv(x,self.basic_conv3_w,self.basic_conv3_b,1) return feat1,feat2############################################################################# class MyYolo: def __init__(self): self.backbone=MyCSPdarknet53_tiny() #conv_forP5 self.conv_forP5_w=None self.conv_forP5_b=None #headP4 self.yolo_headP4_w1=None self.yolo_headP4_w2=None self.yolo_headP4_b1=None self.yolo_headP4_b2=None #headP5 self.yolo_headP5_w1 = None self.yolo_headP5_w2 = None self.yolo_headP5_b1 = None self.yolo_headP5_b2 = None #upsample self.upsample_w=None self.upsample_b=Nonedef forward(self,x): feat1,feat2=self.backbone.forward(x) P5=FusedBasicConv(feat2,self.conv_forP5_w,self.conv_forP5_b,s=1,p=0)#K=1 #out0 out0=FusedBasicConv(P5,self.yolo_headP5_w1,self.yolo_headP5_b1,s=1) out0=F.conv2d(input=out0,weight=self.yolo_headP5_w2,bias=self.yolo_headP5_b2,stride=1,padding=0) #P5_Upsample P5_Upsample=FusedBasicConv(P5,self.upsample_w,self.upsample_b,s=1,p=0)#K=1 P5_Upsample=F.upsample(input=P5_Upsample,scale_factor=2,mode='nearest') #P4 P4=torch.cat([P5_Upsample,feat1],dim=1) #out1 out1=FusedBasicConv(P4,self.yolo_headP4_w1,self.yolo_headP4_b1,s=1) out1=F.conv2d(input=out1,weight=self.yolo_headP4_w2,bias=self.yolo_headP4_b2,stride=1,padding=0) # return out0,out1def load_weight(self,dir): self.backbone.load_weight(dir) #conv_forP5 self.conv_forP5_w=torch.from_numpy(np.fromfile(dir+"\\conv_forP5\\w.bin",dtype=np.float32)).view(256,512,1,1) self.conv_forP5_b=torch.from_numpy(np.fromfile(dir+"\\conv_forP5\\b.bin",dtype=np.float32)) #upsample self.upsample_w=torch.from_numpy(np.fromfile(dir+"\\upsample\\w.bin",dtype=np.float32)).view(128,256,1,1) self.upsample_b=torch.from_numpy(np.fromfile(dir+"\\upsample\\b.bin",dtype=np.float32)) #head4 self.yolo_headP4_w1 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP4\\w1.bin", dtype=np.float32)).view(256,384,3,3) self.yolo_headP4_b1 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP4\\b1.bin", dtype=np.float32)) self.yolo_headP4_w2 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP4\\w2.bin", dtype=np.float32)).view(75,256,1,1) self.yolo_headP4_b2 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP4\\b2.bin", dtype=np.float32)) #head5 self.yolo_headP5_w1 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP5\\w1.bin", dtype=np.float32)).view(512,256,3,3) self.yolo_headP5_b1 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP5\\b1.bin", dtype=np.float32)) self.yolo_headP5_w2 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP5\\w2.bin", dtype=np.float32)).view(75,512,1,1) self.yolo_headP5_b2 = torch.from_numpy(np.fromfile(dir + "\\yolo_headP5\\b2.bin", dtype=np.float32))############################################################################# from yolo import *def backbone_test(): x=torch.randn(1,3,416,416) yolo=YoloBody(anchors_mask=[[3,4,5],[1,2,3]],num_classes=20,phi=0,pretrained=False) yolo.load_state_dict(torch.load("yolov4_tiny_weights_voc.pth")) yolo.eval() mybackbone=MyCSPdarknet53_tiny() mybackbone.load_weight("folded_weights") # f11,f12=mybackbone.forward(x) f21,f22=yolo.backbone.forward(x) print(torch.max(torch.abs(f11-f21))) print(torch.max(torch.abs(f12-f22))) x.numpy().tofile("x.bin") f11.numpy().tofile("feat1.bin") f12.numpy().tofile("feat2.bin")def yolo_test(): #输入 x=torch.randn(10,3,416,416) #baseline yolo=YoloBody(anchors_mask=[[3,4,5], [1,2,3]],num_classes=20,phi=0,pretrained=False) yolo.load_state_dict(torch.load("yolov4_tiny_weights_voc.pth")) yolo.eval() #ours myyolo=MyYolo() myyolo.load_weight("folded_weights") #test o1,o2=myyolo.forward(x) r1,r2=yolo.forward(x) #compare print(torch.max(torch.abs(r1-o1))) print(torch.max(torch.abs(r2-o2)))import os def BasicConvTest(dir,h,w,k,s,p,c,n,i): weight=torch.from_numpy(np.fromfile(dir+"\\w.bin",dtype=np.float32)).view(n,c,k,k) bias=torch.from_numpy(np.fromfile(dir+"\\b.bin",dtype=np.float32)) x=torch.randn(1,c,h,w) out=FusedBasicConv(x=x,w=weight,b=bias,s=s,p=p) if not os.path.exists("BasicConv{}".format(i)): os.mkdir("BasicConv{}".format(i)) weight.numpy().tofile("BasicConv{}\\w.bin".format(i)) bias.numpy().tofile("BasicConv{}\\b.bin".format(i)) x.numpy().tofile("BasicConv{}\\x.bin".format(i)) out.numpy().tofile("BasicConv{}\\out.bin".format(i))def ResBlockTest(dir,h,w,c,n,i): w1=torch.from_numpy(np.fromfile(dir+"\\w1.bin",dtype=np.float32)).view(n,c,3,3) b1=torch.from_numpy(np.fromfile(dir+"\\b1.bin",dtype=np.float32)) w2=torch.from_numpy(np.fromfile(dir+"\\w2.bin",dtype=np.float32)).view(n//2,n//2,3,3) b2=torch.from_numpy(np.fromfile(dir+"\\b2.bin",dtype=np.float32)) w3=torch.from_numpy(np.fromfile(dir+"\\w3.bin",dtype=np.float32)).view(n//2,n//2,3,3) b3=torch.from_numpy(np.fromfile(dir+"\\b3.bin",dtype=np.float32)) w4=torch.from_numpy(np.fromfile(dir+"\\w4.bin",dtype=np.float32)).view(n,n,1,1) b4=torch.from_numpy(np.fromfile(dir+"\\b4.bin",dtype=np.float32)) x=torch.randn(1,c,h,w) #创建保存的文件夹 dir="ResBlock{}".format(i) if not os.path.exists(dir): os.mkdir(dir) y,feat=FusedResBlock(x,n=n,c=c,w1=w1,w2=w2,w3=w3,w4=w4,b1=b1,b2=b2,b3=b3,b4=b4) #保存 x.numpy().tofile(dir+"\\x.bin") y.numpy().tofile(dir+"\\y.bin") feat.numpy().tofile(dir+"\\feat.bin") w1.numpy().tofile(dir+"\\w1.bin") b1.numpy().tofile(dir+"\\b1.bin") w2.numpy().tofile(dir+"\\w2.bin") b2.numpy().tofile(dir+"\\b2.bin") w3.numpy().tofile(dir+"\\w3.bin") b3.numpy().tofile(dir+"\\b3.bin") w4.numpy().tofile(dir+"\\w4.bin") b4.numpy().tofile(dir+"\\b4.bin")if __name__=="__main__": BasicConvTest("folded_weights\\BasicConv1",416,416,3,2,1,3,32,1) BasicConvTest("folded_weights\\BasicConv2",208,208,3,2,1,32,64,2) BasicConvTest("folded_weights\\BasicConv3",13,13,3,1,1,512,512,3) ResBlockTest("folded_weights\\ResBlock1",104,104,64,64,1) ResBlockTest("folded_weights\\ResBlock2",52,52,128,128,2) ResBlockTest("folded_weights\\ResBlock3",26,26,256,256,3) backbone_test() yolo_test()

注:该代码所使用的权重已经经过BN融合操作
效果演示 【python|Yolo4-Tiny代码解读】