YOLO系列解读DAY1—YOLOV1预训练模型

一、说在前面

小伙伴们好,博主很久没有写博客了,略感生疏,不到之处敬请谅解,欢迎指出文中错误,大家一起探讨。欲看视频讲解,可转至博主DouYin、B站,欢迎关注,链接如下:

Github:GitHub - samylee/ExtractionNet: yolov1 extractionnet to pytorch

DY:samylee_csdn

B站:samylee

二、博客系列链接

YOLO系列解读DAY1—YOLOV1预训练模型

YOLO系列解读DAY2—YOLOV1预测代码转换

三、YOLOV1预训练模型

众所周知,在当时预训练模型的性能会直接影响整个检测网络的效果,所以设计出一款好的分类模型是做检测任务的必经之路。分类网络原理部分这里不做进一步阐述了,博主会开设另外系列博客讲解分类网络,敬请期待。

Darknet作者采用自研的extraction网络,在当时已经达到SOTA水平,且速度明显优于同时期SOTA模型,性能如下图所示。

网络设计方面作者采用直到型架构,因当时何恺明等大神尚未提出残差网络,所以直到型还是主流网络架构,如下图所示,其中'M'表示最大池化层,'A'表示平均池化层。

四、Extraction网络转PyTorch

网络转换需要注意三个地方:

1、要求将Extraction网络按照其特定的网络序列复写出PyTorch的网络架构,切不可心浮气躁;

2、Darknet存储参数的序列和PyTorch稍有不同,若遇到BatchNorm架构的卷积层,Darknet会先存储BatchNorm层的参数,进而存储卷积层的参数;

3、图像数据预处理部分,Darknet分类网络的顺序是:BGR2RGB->Norm(1/255)->Resize->Crop,所以我们在复现的时候一定要按照该顺序进行数据预处理操作。

网络复现如下,版本PyTorch1.4,Python3.6,博主暂未加PyTorch的训练代码,因ImageNet训练时间较长,所以直接使用了Extraction的网络模型,若有需要的小伙伴可以评论或私信给博主,博主可添加Extraction分类网络的PyTorch的训练方法。

import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np


class ExtractionNet(nn.Module):
    def __init__(self, in_channels=3, num_classes=1000, init_weights=True, phase='train'):
        super(ExtractionNet, self).__init__()
        self._phase = phase
        self.features = self.make_layers(in_channels=in_channels, out_channels=num_classes)
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        for feature in self.features:
            x = feature(x)

        if self._phase == 'test':
            x = torch.flatten(x, 1)
            x = F.softmax(x, dim=-1)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

    def make_layers(self, in_channels=3, out_channels=1000):
        # conv: out_channels, kernel_size, stride, batchnorm, activate
        # maxpool: kernel_size stride
        params = [
            [64, 7, 2, True, 'leaky'],
            ['M', 2, 2],
            [192, 3, 1, True, 'leaky'],
            ['M', 2, 2],
            [128, 1, 1, True, 'leaky'],
            [256, 3, 1, True, 'leaky'],
            [256, 1, 1, True, 'leaky'],
            [512, 3, 1, True, 'leaky'],
            ['M', 2, 2],
            [256, 1, 1, True, 'leaky'],
            [512, 3, 1, True, 'leaky'],
            [256, 1, 1, True, 'leaky'],
            [512, 3, 1, True, 'leaky'],
            [256, 1, 1, True, 'leaky'],
            [512, 3, 1, True, 'leaky'],
            [256, 1, 1, True, 'leaky'],
            [512, 3, 1, True, 'leaky'],
            [512, 1, 1, True, 'leaky'],
            [1024, 3, 1, True, 'leaky'],
            ['M', 2, 2],
            [512, 1, 1, True, 'leaky'],
            [1024, 3, 1, True, 'leaky'],
            [512, 1, 1, True, 'leaky'],
            [1024, 3, 1, True, 'leaky'],
            [out_channels, 1, 1, False, 'leaky'], # classifier
            ['A']
        ]

        module_list = nn.ModuleList()
        for i, v in enumerate(params):
            modules = nn.Sequential()
            if v[0] == 'M':
                modules.add_module(f"maxpool_{i}", nn.MaxPool2d(kernel_size=v[1], stride=v[2], padding=int((v[1] - 1) // 2)))
            elif v[0] == 'A':
                modules.add_module(f"avgpool_{i}", nn.AdaptiveAvgPool2d((1, 1)))
            else:
                modules.add_module(
                    f"conv_{i}",
                    nn.Conv2d(
                        in_channels,
                        v[0],
                        kernel_size=v[1],
                        stride=v[2],
                        padding=(v[1] - 1) // 2,
                        bias=not v[3]
                    )
                )
                if v[3]:
                    modules.add_module(f"bn_{i}", nn.BatchNorm2d(v[0]))
                modules.add_module(f"act_{i}", nn.LeakyReLU(0.1) if v[4] == 'leaky' else nn.ReLU())
                in_channels = v[0]
            module_list.append(modules)
        return module_list


def load_darknet_weights(model, weights_path):
    # Open the weights file
    with open(weights_path, "rb") as f:
        # First five are header values
        header = np.fromfile(f, dtype=np.int32, count=4)
        header_info = header  # Needed to write header when saving weights
        seen = header[3]  # number of images seen during training
        weights = np.fromfile(f, dtype=np.float32)  # The rest are weights

    ptr = 0
    for module in model.features:
        if isinstance(module[0], nn.Conv2d):
            conv_layer = module[0]
            if isinstance(module[1], nn.BatchNorm2d):
                # Load BN bias, weights, running mean and running variance
                bn_layer = module[1]
                num_b = bn_layer.bias.numel()  # Number of biases
                # Bias
                bn_b = torch.from_numpy(
                    weights[ptr: ptr + num_b]).view_as(bn_layer.bias)
                bn_layer.bias.data.copy_(bn_b)
                ptr += num_b
                # Weight
                bn_w = torch.from_numpy(
                    weights[ptr: ptr + num_b]).view_as(bn_layer.weight)
                bn_layer.weight.data.copy_(bn_w)
                ptr += num_b
                # Running Mean
                bn_rm = torch.from_numpy(
                    weights[ptr: ptr + num_b]).view_as(bn_layer.running_mean)
                bn_layer.running_mean.data.copy_(bn_rm)
                ptr += num_b
                # Running Var
                bn_rv = torch.from_numpy(
                    weights[ptr: ptr + num_b]).view_as(bn_layer.running_var)
                bn_layer.running_var.data.copy_(bn_rv)
                ptr += num_b
            else:
                # Load conv. bias
                num_b = conv_layer.bias.numel()
                conv_b = torch.from_numpy(weights[ptr: ptr + num_b]).view_as(conv_layer.bias)
                conv_layer.bias.data.copy_(conv_b)
                ptr += num_b
            # Load conv. weights
            num_w = conv_layer.weight.numel()
            conv_w = torch.from_numpy(weights[ptr: ptr + num_w]).view_as(conv_layer.weight)
            conv_layer.weight.data.copy_(conv_w)
            ptr += num_w


if __name__ == '__main__':
    # load moel
    checkpoint_path = "extraction.weights"
    model = ExtractionNet(phase='test')
    load_darknet_weights(model, checkpoint_path)
    model.eval()

    # model input size
    net_size = 224

    # load img
    img = cv2.imread('dog.jpg')
    # img bgr2rgb
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # resize img
    h, w, _ = img_rgb.shape
    new_h, new_w = h, w
    if w < h:
        new_h = (h * net_size) // w
        new_w = net_size
    else:
        new_w = (w * net_size) // h
        new_h = net_size
    img_resize = cv2.resize(img_rgb, (new_w, new_h))

    # crop img
    cut_w = (new_w - net_size) // 2
    cut_h = (new_h - net_size) // 2
    img_crop = img_resize[cut_h:cut_h + net_size, cut_w:cut_w + net_size, :]

    # float
    img_crop = torch.from_numpy(img_crop.transpose((2, 0, 1)))
    img_float = img_crop.float().div(255).unsqueeze(0)

    # forward
    result = model(img_float)

    print(result.topk(5))

五、写在后面

小伙伴们若能坚持完成YOLO系列的代码和原理学习,相信能对图像检测任务有个全新的认识,跟随博主的脚步,培养自己的动手能力吧!希望博主也能坚持将该系列做下去,一起加油!!!

六、参考

1、ImageNet Classification

2、https://github.com/AlexeyAB/darknet

  • 1
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 6
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值