zl程序教程

您现在的位置是:首页 >  后端

当前栏目

【图像分类】用通俗易懂代码的复现EfficientNetV2,入门的绝佳选择(pytorch)

入门PyTorch代码 选择 通俗易懂 图像 分类 复现
2023-09-14 09:05:43 时间

目录

摘要

代码实现

激活函数

SE模块

 定义MBConv模块和Fused-MBConv模块

主体模块

完整代码


摘要

上周学习了EfficientNetV2的论文,并对其进行了翻译,如果对论文感兴趣的可以参考我的文章:https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/117399085

在EfficientNets的第一个版本中存在三个缺点:

(1) 用非常大的图像尺寸训练很慢;而且输入较大的图像必须以较小的批大小训练这些模型,这大大减慢了训练速度,但是精度反而下降了。

(2) 在网络浅层中使用Depthwise convolutions速度会很慢。

(3) 每个阶段都按比例放大是次优的。 

EfficientNetV2针对这三个方面的缺点做了改进:

1、针对图像的大小的问题,作者提出了自适应正则化的渐进式学习的方法,详见论文的4.2节

2、针对Depthwise convolutions在早期层中很慢的问题,作者提出了 Fused-MBConv 模块来代替部分的MBConv。

3、针对每个阶段都按比例放大是次优的问题,作者使用非均匀缩放策略来逐步添加 到后期阶段。

作者对EfficientNetV2做了三方面的总结

          •我们引入了 EfficientNetV2,这是一个新的更小、更快的模型系列。 通过我们的训练感知NAS 和扩展发现,EfficientNetV2 在训练速度和参数效率方面都优于以前的模型。

          • 我们提出了一种改进的渐进式学习方法,它可以根据图像大小自适应地调整正则化。 我们表明它可以加快训练速度,同时提高准确性。

          • 我们在 ImageNet、CIFAR、Cars 和 Flowers 数据集上展示了比现有技术快 11 倍的训练速度和 6.8 倍的参数效率。

总之,一句话,我们的新模型又快又准而且还小,大家赶快用吧!下面我就讲讲如何使用Pytorch实现EfficientNetV2。

代码实现

EfficientNetV2和EfficientNet一样也是一个家族模型,包括:efficientnetv2_s、efficientnetv2_m,、efficientnetv2_l、efficientnetv2_xl。所以我们要实现四个模型。

激活函数

激活函数使用SiLU激活函数,我对激活函数做了总结,感兴趣的可以查看:CNN基础——激活函数_AI浩-CSDN博客

# SiLU (Swish) activation function
if hasattr(nn, 'SiLU'):
    SiLU = nn.SiLU
else:
    # For compatibility with old PyTorch versions
    class SiLU(nn.Module):
        def forward(self, x):
            return x * torch.sigmoid(x)

SE模块

SE模块,我在前面的文章中已经介绍了。现在直接将SE模块拿过来使用。

class SELayer(nn.Module):
    def __init__(self, inp, oup, reduction=4):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
                nn.Linear(oup, _make_divisible(inp // reduction, 8)),
                SiLU(),
                nn.Linear(_make_divisible(inp // reduction, 8), oup),
                nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y

 定义MBConv模块和Fused-MBConv模块

这两个模块是整个模型实现的核心,模块的详细构造如下图:

我们可以看到MBConv模块,经过1×1的卷积,然后channel放大四倍,再经过depthwise conv3×3的卷积,然后经过SE模块后,再经过1×1的卷积,把channel恢复到输入的大小,最后和上层的输入融合。

Fused-MBConv模块比MBConv模块简单些,先经过3×3的卷积,把channel放大四倍,然后经过SE模块,再经过1×1的卷积,最后和上层的输入融合。下面是实现MBConv模块和Fused-MBConv模块的详细代码:

class MBConv(nn.Module):
    """
     定义MBConv模块和Fused-MBConv模块,将fused设置为1或True是Fused-MBConv,否则是MBConv
    :param inp:输入的channel
    :param oup:输出的channel
    :param stride:步长,设置为1时图片的大小不变,设置为2时,图片的面积变为原来的四分之一
    :param expand_ratio:放大的倍率
    :return:
    """
    def __init__(self, inp, oup, stride, expand_ratio, fused):
        super(MBConv, self).__init__()
        assert stride in [1, 2]
        hidden_dim = round(inp * expand_ratio)
        self.identity = stride == 1 and inp == oup
        if fused:
            self.conv = nn.Sequential(
                # fused
                nn.Conv2d(inp, hidden_dim, 3, stride, 1, bias=False),
                nn.BatchNorm2d(hidden_dim),
                SiLU(),
                SELayer(inp, hidden_dim),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:

            self.conv = nn.Sequential(
                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                SiLU(),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                SiLU(),
                SELayer(inp, hidden_dim),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.identity:
            return x + self.conv(x)
        else:
            return self.conv(x)

主体模块

class EfficientNetv2(nn.Module):
    def __init__(self, cfgs, num_classes=1000, width_mult=1.):
        super(EfficientNetv2, self).__init__()
        self.cfgs = cfgs

        # building first layer
        input_channel = _make_divisible(24 * width_mult, 8)
        layers = [conv_3x3_bn(3, input_channel, 2)]
        # building inverted residual blocks
        block = MBConv
        for t, c, n, s, fused in self.cfgs:
            output_channel = _make_divisible(c * width_mult, 8)
            for i in range(n):
                layers.append(block(input_channel, output_channel, s if i == 0 else 1, t, fused))
                input_channel = output_channel
        self.features = nn.Sequential(*layers)
        # building last several layers
        output_channel = _make_divisible(1792 * width_mult, 8) if width_mult > 1.0 else 1792
        self.conv = conv_1x1_bn(input_channel, output_channel)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(output_channel, num_classes)

        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.conv(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        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))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.001)
                m.bias.data.zero_()

 理解这段代码,我们还需要了解输入参数cfgs,以efficientnetv2_s为例:

def efficientnetv2_s(**kwargs):
    """
    Constructs a EfficientNetV2-S model
    """
    cfgs = [
        # t, c, n, s, fused
        [1,  24,  2, 1, 1],
        [4,  48,  4, 2, 1],
        [4,  64,  4, 2, 1],
        [4, 128,  6, 2, 0],
        [6, 160,  9, 1, 0],
        [6, 272, 15, 2, 0],
    ]
    return EfficientNetv2(cfgs, **kwargs)

 第一列“t”指的是MBConv模块和Fused-MBConv模块第一个输入后放大的倍率。

 第二列“c”,channel,指的是输出的channel。

第三列“n”,指定的是MBConv模块和Fused-MBConv模块堆叠的个数。

第四列“s”,指的是卷积的步长,步长为1,图片的大小不变,步长为图片的面积缩小为原来的四分之一,实现降维。

第五列“fused”,选择MBConv模块或Fused-MBConv模块,为1这是Fused-MBConv模块,0则是MBConv模块,对应了前面摘要提过了,在浅层用Fused-MBConv代替MBConv。

完整代码

import torch
import torch.nn as nn
import math

__all__ = ['efficientnetv2_s', 'efficientnetv2_m', 'efficientnetv2_l', 'efficientnetv2_xl']


from torchsummary import summary

#这个函数的目的是确保Channel能被8整除。
def _make_divisible(v, divisor, min_value=None):
    """
    这个函数的目的是确保Channel能被8整除。
    :param v:
    :param divisor:
    :param min_value:
    :return:
    """
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


# SiLU (Swish) activation function
if hasattr(nn, 'SiLU'):
    SiLU = nn.SiLU
else:
    # For compatibility with old PyTorch versions
    class SiLU(nn.Module):
        def forward(self, x):
            return x * torch.sigmoid(x)

 
class SELayer(nn.Module):
    def __init__(self, inp, oup, reduction=4):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
                nn.Linear(oup, _make_divisible(inp // reduction, 8)),
                SiLU(),
                nn.Linear(_make_divisible(inp // reduction, 8), oup),
                nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


def conv_3x3_bn(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
        nn.BatchNorm2d(oup),
        SiLU()
    )


def conv_1x1_bn(inp, oup):
    return nn.Sequential(
        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        SiLU()
    )


class MBConv(nn.Module):
    """
     定义MBConv模块和Fused-MBConv模块,将fused设置为1或True是Fused-MBConv,否则是MBConv
    :param inp:输入的channel
    :param oup:输出的channel
    :param stride:步长,设置为1时图片的大小不变,设置为2时,图片的面积变为原来的四分之一
    :param expand_ratio:放大的倍率
    :return:
    """
    def __init__(self, inp, oup, stride, expand_ratio, fused):
        super(MBConv, self).__init__()
        assert stride in [1, 2]
        hidden_dim = round(inp * expand_ratio)
        self.identity = stride == 1 and inp == oup
        if fused:
            self.conv = nn.Sequential(
                # fused
                nn.Conv2d(inp, hidden_dim, 3, stride, 1, bias=False),
                nn.BatchNorm2d(hidden_dim),
                SiLU(),
                SELayer(inp, hidden_dim),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )
        else:

            self.conv = nn.Sequential(
                # pw
                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                SiLU(),
                # dw
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                SiLU(),
                SELayer(inp, hidden_dim),
                # pw-linear
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            )

    def forward(self, x):
        if self.identity:
            return x + self.conv(x)
        else:
            return self.conv(x)


class EfficientNetv2(nn.Module):
    def __init__(self, cfgs, num_classes=1000, width_mult=1.):
        super(EfficientNetv2, self).__init__()
        self.cfgs = cfgs

        # building first layer
        input_channel = _make_divisible(24 * width_mult, 8)
        
        layers = [conv_3x3_bn(3, input_channel, 2)]
        # building inverted residual blocks
        block = MBConv
        for t, c, n, s, fused in self.cfgs:
            output_channel = _make_divisible(c * width_mult, 8)
            for i in range(n):
                layers.append(block(input_channel, output_channel, s if i == 0 else 1, t, fused))
                input_channel = output_channel
        self.features = nn.Sequential(*layers)
        # building last several layers
        output_channel = _make_divisible(1792 * width_mult, 8) if width_mult > 1.0 else 1792
        self.conv = conv_1x1_bn(input_channel, output_channel)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(output_channel, num_classes)

        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.conv(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        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))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.001)
                m.bias.data.zero_()


def efficientnetv2_s(**kwargs):
    """
    Constructs a EfficientNetV2-S model
    """
    cfgs = [
        # t, c, n, s, fused
        [1,  24,  2, 1, 1],
        [4,  48,  4, 2, 1],
        [4,  64,  4, 2, 1],
        [4, 128,  6, 2, 0],
        [6, 160,  9, 1, 0],
        [6, 272, 15, 2, 0],
    ]
    return EfficientNetv2(cfgs, **kwargs)


def efficientnetv2_m(**kwargs):
    """
    Constructs a EfficientNetV2-M model
    """
    cfgs = [
        # t, c, n, s, fused
        [1,  24,  3, 1, 1],
        [4,  48,  5, 2, 1],
        [4,  80,  5, 2, 1],
        [4, 160,  7, 2, 0],
        [6, 176, 14, 1, 0],
        [6, 304, 18, 2, 0],
        [6, 512,  5, 1, 0],
    ]
    return EfficientNetv2(cfgs, **kwargs)


def efficientnetv2_l(**kwargs):
    """
    Constructs a EfficientNetV2-L model
    """
    cfgs = [
        # t, c, n, s, fused
        [1,  32,  4, 1, 1],
        [4,  64,  7, 2, 1],
        [4,  96,  7, 2, 1],
        [4, 192, 10, 2, 0],
        [6, 224, 19, 1, 0],
        [6, 384, 25, 2, 0],
        [6, 640,  7, 1, 0],
    ]
    return EfficientNetv2(cfgs, **kwargs)


def efficientnetv2_xl(**kwargs):
    """
    Constructs a EfficientNetV2-XL model
    """
    cfgs = [
        # t, c, n, s, fused
        [1,  32,  4, 1, 1],
        [4,  64,  8, 2, 1],
        [4,  96,  8, 2, 1],
        [4, 192, 16, 2, 0],
        [6, 256, 24, 1, 0],
        [6, 512, 32, 2, 0],
        [6, 640,  8, 1, 0],
    ]
    return EfficientNetv2(cfgs, **kwargs)

if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = efficientnetv2_s()
    model.to(device)
    summary(model, (3, 224, 224))