ECA-NET代码学习（pytorch）

文件目录：

• ECANet-master
  • models
    • eca_mobilenetv2.py
    • eca_module.py
    • eca_ns.py
    • eca_resnet.py
• main.py
• paras_flops.py

ECA_module

文中注释部分

import torch
from torch import nn
from torch.nn.parameter import Parameter

# ECA模块
class eca_layer(nn.Module):
    """Constructs a ECA module.

    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """
    def __init__(self, channel, k_size=3):
        super(eca_layer, self).__init__() # super类的作用是继承的时候，调用含super的哥哥的基类__init__函数。
        self.avg_pool = nn.AdaptiveAvgPool2d(1) # 全局平均池化
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size,
                              padding=(k_size - 1) // 2, bias=False) # 一维卷积
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x: input features with shape [b, c, h, w]
        b, c, h, w = x.size() # b代表b个样本，c为通道数，h为高度，w为宽度

        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        # torch.squeeze()这个函数主要对数据的维度进行压缩,torch.unsqueeze()这个函数 主要是对数据维度进行扩充
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion多尺度信息融合
        y = self.sigmoid(y)
        # 原网络中克罗内克积，也叫张量积，为两个任意大小矩阵间的运算
        return x * y.expand_as(x)

torch.squeeze()和 unsqueeze()

数据维度扩充unsqueeze()

对数据维度进行扩充。a.squeeze(N) 就是在a中指定位置N加上一个维数为1的维度。

import torch

a = torch.randn(1, 3)
print(a)
print(a.shape)
# 在位置1给a增加维度为1的维度
b = torch.unsqueeze(a, 1)
c = a.unsqueeze(0)
print(b)
print(b.shape)
print(c)
print(c.shape)

结果：

数据维度压缩squeeze()

对数据维度进行压缩。

d = torch.squeeze(c)
print(d)
print(d.shape)

结果：

torch.squeeze()和torch.unsqueeze()

import torch
a = torch.randn(1, 3)
print("ordinary")
print(a)
b = a.squeeze(-1).transpose(-1, -2)
print("squeeze")
print(b)
c = (a.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
print("unsqueeze")
print(c)

结果

ECA-ResNet

import torch.nn as nn
import math
# import torch.utils.model_zoo as model_zoo
from .eca_module import eca_layer

# 3*3卷积
def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class ECABasicBlock(nn.Module):

    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, k_size=3):
        super(ECABasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes, 1)
        self.bn2 = nn.BatchNorm2d(planes)
        self.eca = eca_layer(planes, k_size)
        self.downsample = downsample
        # 步长
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.eca(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ECABottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, k_size=3):
        super(ECABottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.eca = eca_layer(planes * 4, k_size)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        out = self.eca(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

torch.view()

可以调整tensor的形状，元素个数保持不变。

import torch as t

a = t.arange(0, 6)
print(a)
b = a.view(2, 3)
print(b)

结果：

torch.resize()

import torch as t

a = t.arange(0, 6)
print(a)
b = a.view(2, 3)
print(b)

c = b.resize_(1, 3)
print(c)
d = b.resize_(3, 3)
print(d)

结果：

nn.Sequential()

nn.Sequential()一个有序的容器，神经网络模块将按照在传入构造器顺序依次被添加到计算图中执行，同时以神经网络模块为元素的有序字典也可以作为传入参数。

self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
            )