手动实现yolov3-tiny模型结构

最新推荐文章于 2024-05-02 01:00:00 发布

朱小丰

最新推荐文章于 2024-05-02 01:00:00 发布

阅读量1.3k

点赞数 3

分类专栏：检测文章标签： yolov3-tiny

本文链接：https://blog.csdn.net/weixin_41866216/article/details/108406384

版权

检测专栏收录该内容

6 篇文章 0 订阅

订阅专栏

preview

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
----------------------------------------------------------------

朱小丰

关注

3
点赞
踩
14

收藏

觉得还不错? 一键收藏
0
评论
手动实现yolov3-tiny模型结构

import torchimport torch.nn as nnimport numpy as npfrom torchsummary import summaryimport torch.nn.functional as Fclass convBatchReluBlock(nn.Module): def __init__(self, in_c, out_c, k, s, p): super().__init__() self.cblr = nn...
复制链接

扫一扫