今天搭了個(gè)yolo v5的網(wǎng)絡(luò)模型，不知道訓(xùn)練效果怎么樣，反正模型是搭起來了，反正又不要求百分百?gòu)?fù)現(xiàn)，嘿嘿嘿嘿嘿。

1.主干網(wǎng)絡(luò)：cspdarknet.py

import torch import torch.nn as nn from torchsummary import summarydef autopad(k, p=None):if p is None:p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-padreturn pclass CBLx1(nn.Module):def __init__(self, ch_in, ch_out, k=3, s=2, p=None, groups=1):super(CBLx1, self).__init__()self.conv = nn.Conv2d(ch_in, ch_out, k, s, autopad(k, p), groups=groups, bias=False)self.bn = nn.BatchNorm2d(ch_out)self.act = nn.SiLU()def __call__(self, x):x = self.conv(x)x = self.bn(x)return self.act(x)class CBL(nn.Module):def __init__(self, ch_in, ch_out, k=3, s=2, p=None, groups=1):super(CBL, self).__init__()self.conv = nn.Conv2d(ch_in, ch_out, k, s, autopad(k, p), groups=groups, bias=False)self.bn = nn.BatchNorm2d(ch_out)self.act = nn.LeakyReLU()def __call__(self, x):x = self.conv(x)x = self.bn(x)return self.act(x)class Focus(nn.Module):def __init__(self, ch_in, ch_out, k=1, s=1, p=None, groups=1):super(Focus, self).__init__()self.cbl = CBLx1(ch_in*4, ch_out, k, s, autopad(k,p), groups)def __call__(self, x):out = torch.cat([x[..., ::2, ::2], # 從第0行第0列開始，每隔一個(gè)元素取值x[..., 1::2, ::2], # 從第1行第0列開始，每隔一個(gè)元素取值x[..., ::2, 1::2], # 從第0行第1列開始，每隔一個(gè)元素取值x[..., 1::2, 1::2] # 從第1行第1列開始，每隔一個(gè)元素取值], 1)return self.cbl(out)class Bottleneck(nn.Module):def __init__(self, ch_in, ch_out, blocks, shortcut=True):super(Bottleneck, self).__init__()self.add = (shortcut and ch_in == ch_out)hidden_chanel = ch_out // 2unit_layer = nn.Sequential(CBL(ch_in, ch_out, k=1, s=1),CBL(2 * hidden_chanel, ch_out, 3, 1))self.unit_layer_n = nn.Sequential(*[unit_layer for _ in range(blocks)])def __call__(self, x):if self.add:return x + self.unit_layer_n(x)else:return self.unit_layer_n(x)class CSP1_n(nn.Module):def __init__(self, ch_in, ch_out, k=1, s=1, p=None, groups=1, n=1):super(CSP1_n, self).__init__()c_ = int(ch_out // 2)self.up = nn.Sequential(CBL(ch_in, c_, k, s, autopad(k, p), groups),Bottleneck(c_, c_, n),nn.Conv2d(c_, c_, 1))self.bottom = nn.Conv2d(ch_in, c_, 1, 1, 0)self.tie = nn.Sequential(nn.BatchNorm2d(c_ * 2),nn.LeakyReLU(),nn.Conv2d(c_ * 2, ch_out, 1, 1, 0, bias=False))def __call__(self, x):total = torch.cat([self.up(x), self.bottom(x)], dim=1)out = self.tie(total)return outclass CSP2_n(nn.Module):def __init__(self, ch_in, ch_out, k=1, s=1, p=None, groups=1, n=1):super(CSP2_n, self).__init__()c_ = ch_out // 2self.up = nn.Sequential(CBL(ch_in, c_, k, s, autopad(k, p), groups),Bottleneck(c_, c_, n),nn.Conv2d(c_, ch_out, 1),)self.tie = nn.Sequential(nn.BatchNorm2d(ch_out),nn.LeakyReLU(),nn.Conv2d(ch_out, ch_out, 1, 1, 0, bias=False))def __call__(self, x):total = self.up(x)out = self.tie(total)return outclass SPP(nn.Module):def __init__(self, ch_in, ch_out, k=(5, 9, 13)):super(SPP, self).__init__()hiddel_channel = ch_in // 2length = len(k) + 1self.conv1 = CBL(ch_in, hiddel_channel, 1, 1)self.max_pool = nn.ModuleList(nn.MaxPool2d(kernel_size=x, stride=1, padding=x//2) for x in k)self.conv2 = CBL(hiddel_channel * length, ch_out, 1, 1)def forward(self, x):x = self.conv1(x)out = torch.cat([x] + [m(x) for m in self.max_pool], 1)out = self.conv2(out)return outclass CSPDarkNet(nn.Module):def __init__(self, gd=0.33, gw=0.5):super(CSPDarkNet, self).__init__()self.truck_big = nn.Sequential(Focus(3, 32),CBL(32, 64, k=3, s=2, p=1),CSP1_n(64, 64, n=3),CBL(64, 128, k=3, s=2, p=1),CSP1_n(128, 128, n=3),)self.truck_middle = nn.Sequential(CBL(128, 256, k=3, s=2, p=1),CSP1_n(256, 256, n=3),)self.truck_small = nn.Sequential(CBL(256, 512, k=3, s=2, p=1),SPP(512, 512))def forward(self, x):h_big = self.truck_big(x) # (80,80)h_middle = self.truck_middle(h_big) # (40,40)h_small = self.truck_small(h_middle) # (20,20)return h_big, h_middle, h_smallif __name__ == '__main__':model = CSPDarkNet(gd=0.33, gw=0.5)summary(model, input_size=(3, 640, 640), batch_size=1, device="cpu")

2.yolo.py

import torch import torch.nn as nn from cspdarknet import * from torch.nn import Upsample as UpSampleclass YOLO(nn.Module):def __init__(self, nc=80):super(YOLO, self).__init__()self.nc = ncself.backone = CSPDarkNet()self.neck_small = nn.Sequential(CSP2_n(512, 512, n=3),CBL(512, 256, 1, 1, 0))self.up_middle = nn.Sequential(UpSample(scale_factor=2))self.out_set_middle = nn.Sequential(CSP2_n(512, 256, n=3),CBL(256, 128, 1, 1, 0),)self.up_big = nn.Sequential(UpSample(scale_factor=2))self.out_set_tie_big = nn.Sequential(CSP2_n(256, 128, n=3))self.pan_middle = nn.Sequential(CBL(128, 128, 3, 2, 1))self.out_set_tie_middle = nn.Sequential(CSP2_n(256, 256, n=3))self.pan_small = nn.Sequential(CBL(256, 256, 3, 2, 1))self.out_set_tie_small = nn.Sequential(CSP2_n(512, 512, n=3))# ------------------------------Prediction--------------------------------# predictionbig_ = round(128)middle = round(256)small_ = round(512)self.out_big = nn.Sequential(nn.Conv2d(big_, 3 * (5 + nc), 1, 1, 0))self.out_middle = nn.Sequential(nn.Conv2d(middle, 3 * (5 + nc), 1, 1, 0))self.out_small = nn.Sequential(nn.Conv2d(small_, 3 * (5 + nc), 1, 1, 0))def __call__(self, x):big, middle, small = self.backone(x)# 第一部分：# 1.CSP2_n + CBLneck_small = self.neck_small(small)# 2.從上采樣到CBLup_middle = self.up_middle(neck_small)middle_cat = torch.cat([up_middle, middle], dim=1)out_set_middle = self.out_set_middle(middle_cat)# 3.上采樣到CSP2_nup_big = self.up_big(out_set_middle) # torch.Size([2, 128, 76, 76])big_cat = torch.cat([up_big, big], dim=1)out_set_tie_big = self.out_set_tie_big(big_cat)# 4.后面out_big = self.out_big(out_set_tie_big)out_big = out_big.view(-1, 3, (5+self.nc), out_big.shape[-2], out_big.shape[-1])out_big = out_big.permute(0, 1, 3, 4, 2)# 第二部分：# 1.CBL-7到CSP-7neck_tie_middle = torch.cat([self.pan_middle(out_set_tie_big), out_set_middle], dim=1)up_middle = self.out_set_tie_middle(neck_tie_middle)# 2.后面out_middle = self.out_middle(up_middle)out_middle = out_middle.view(-1, 3, (5+self.nc), out_middle.shape[-2], out_middle.shape[-1])out_middle = out_middle.permute(0, 1, 3, 4, 2)# 第三部分：# 1.CBL-8到CSP-8neck_tie_small = torch.cat([self.pan_small(up_middle), neck_small], dim=1)out_set_small = self.out_set_tie_small(neck_tie_small)# 2.后面部分out_small = self.out_small(out_set_small)out_small = out_small.view(-1, 3, (5 + self.nc), out_small.shape[-2], out_small.shape[-1])out_small = out_small.permute(0, 1, 3, 4, 2)return out_small, out_middle, out_bigif __name__ == '__main__':net = YOLO()# from torchsummary import summary# summary(net, input_size=(3, 640, 640), device="cpu")a = torch.randn(2, 3, 640, 640)y = net(a)print(y[0].shape, y[1].shape, y[2].shape)

總結(jié)

以上是生活随笔為你收集整理的（一）yolo5的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。