【YOLO v1】模型搭建 ｜ model ｜ 代码

YOLO V1 模型

import torch
import torch.nn as nn
from torchsummary import summary


def build_block(in_channel, out_channel, kernel_size, stride=1, maxpool=False):
    padding = kernel_size//2
    block = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
                          nn.BatchNorm2d(out_channel),
                          nn.LeakyReLU(0.1, inplace=True))
    if maxpool:
        block.add_module("maxpool2d", nn.MaxPool2d(kernel_size=2, stride=2))
    return block


class YOLOv1(nn.Module):
    def __init__(self, num_classes=20):
        super(YOLOv1, self).__init__()
        self.num_classes = num_classes
        self.layer1 = nn.Sequential(build_block(3, 64, 7, stride=2, maxpool=True))
        self.layer2 = nn.Sequential(build_block(64, 192, 3, maxpool=True))

        self.layer3 = nn.Sequential(build_block(192, 128, 1),
                                    build_block(128, 256, 3),
                                    build_block(256, 256, 1),
                                    build_block(256, 512, 3, maxpool=True))

        self.layer4 = nn.Sequential(build_block(512, 256, 1),
                                    build_block(256, 512, 3),
                                    build_block(512, 256, 1),
                                    build_block(256, 512, 3),
                                    build_block(512, 256, 1),
                                    build_block(256, 512, 3),
                                    build_block(512, 256, 1),
                                    build_block(256, 512, 3),
                                    build_block(512, 512, 1),
                                    build_block(512, 1024, 3, maxpool=True))

        self.layer5 = nn.Sequential(build_block(1024, 512, 1),
                                    build_block(512, 1024, 3),
                                    build_block(1024, 512, 1),
                                    build_block(512, 1024, 3),
                                    build_block(1024, 1024, 3),
                                    build_block(1024, 1024, 3, stride=2))

        self.layer6 = nn.Sequential(build_block(1024, 1024, 3),
                                    build_block(1024, 1024, 3))

        self.layer7 = nn.Sequential(nn.Flatten(),
                                    nn.Linear(7 * 7 * 1024, 4096, bias=True),
                                    nn.LeakyReLU(0.1),
                                    nn.Dropout(p=0.5),
                                    nn.Linear(4096, 7 * 7 * (self.num_classes + 5 * 2), bias=True),
                                    nn.Sigmoid(),
                                    )

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.layer6(x)
        x = self.layer7(x)
        x = x.view(-1, 7, 7, self.num_classes + 5 * 2)
        return x


if __name__ == '__main__':
    net = YOLOv1()
    print(summary(net, (3, 448, 448)))

    x = torch.randn((1, 3, 448, 448))
    print(net(x).shape)