ASPP 代码解析 pytorch

无妄无望

2843人浏览 · 2023-10-14 12:49:41

无妄无望 · 2023-10-14 12:49:41 发布

1.结构图

如图所示，ASPP 本质上由一个1×1的卷积（最左侧绿色） + 池化金字塔（中间三个蓝色） + ASPP Pooling（最右侧三层）组成。而池化金字塔各层的膨胀因子可自定义，从而实现自由的多尺度特征提取。

2.代码讲解

ASPP 是一种用于图像语义分割的技术，通过在不同的空洞率下进行卷积操作，从而捕捉不同尺度的上下文信息。

ASPPConv

class ASPPConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, dilation):
        modules = [
            nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        ]
        super(ASPPConv, self).__init__(*modules)

这段代码是一个 ASPPConv 类，继承自 nn.Sequential，用于实现 Atrous Spatial Pyramid Pooling (ASPP) 中的卷积操作。该类的输入参数包括

in_channels : 输入通道数

out_channels : 输出通道数

dilation : 空洞率

主要内容在modules中，包含一个卷积层、一个批归一化层和一个ReLU激活函数

其中super(ASPPConv, self).__init__(*modules)是一个 Python 中的类的构造函数，ASPPConv 是类名，*modules 表示接受任意数量的参数并将它们作为元组传递给函数。super() 函数返回一个临时对象，该对象继承了父类的方法。在这里，super(ASPPConv, self) 表示调用 ASPPConv 的父类的构造函数，并将 self 作为参数传递给它。

简单来说，调用ASPPConv时往里填输入通道数，输出通道数和空洞率三个数据就行

示例，A = ASPPConv（in_channels =224，out_channels=224，dilation=6）

ASPPPooling

class ASPPPooling(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(ASPPPooling, self).__init__(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU())
 
    def forward(self, x):
        size = x.shape[-2:]
        x = super(ASPPPooling, self).forward(x)
        return F.interpolate(x, size=size, mode='bilinear', align_corners=False)

这是一个定义了ASPP模块中的池化层的类，该类接受两个参数：

in_channels ：输入通道数

out_channels ：输出通道数

具体来说，包含了一个自适应平均池化层、一个 1x1 卷积层、一个批归一化层和一个 ReLU 激活函数层。

自适应平均池化层的作用是将输入的特征图进行平均池化，输出一个大小为 1x1 的特征图。

1x1 卷积层的作用是对输入的特征图进行卷积操作，输出一个新的特征图。

批归一化层的作用是对卷积输出的特征图进行归一化处理，加速模型训练。

ReLU 激活函数层则是对特征图进行非线性变换，增强模型的表达能力。

在前向传播过程中

“size = x.shape[-2:]”表示获取一个张量 x 的最后两个维度的大小，也就是 x 的高度和宽度

“x = super(ASPPPooling, self).forward(x)”表示调用父类 ASPPPooling 的 forward 方法，并将参数 x 传递给该方法进行处理

“F.interpolate(x, size=size, mode='bilinear', align_corners=False)”表示对池化后的x进行双线性插值将特征图的大小还原回原始大小

总的来说，在ASPPPooling中，先是池化富集信息，再通过插值恢复图像尺度

ASPP

class ASPP(nn.Module):
    def __init__(self, in_channels, atrous_rates):
        super(ASPP, self).__init__()
        out_channels = 256
        modules = []
        modules.append(nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()))
 
        rate1, rate2, rate3 = tuple(atrous_rates)
        modules.append(ASPPConv(in_channels, out_channels, rate1))
        modules.append(ASPPConv(in_channels, out_channels, rate2))
        modules.append(ASPPConv(in_channels, out_channels, rate3))
        modules.append(ASPPPooling(in_channels, out_channels))
 
        self.convs = nn.ModuleList(modules)
 
        self.project = nn.Sequential(
            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(0.5))
 
    def forward(self, x):
        res = []
        for conv in self.convs:
            res.append(conv(x))
        res = torch.cat(res, dim=1)
        return self.project(res)

这是一个定义了模块ASPP的类，该类接受两个参数：

in_channels : 输入通道数

dilation : 空洞率

out_channels输出通道数的值在给的代码中被钉死。

ASPP模块包含一个1x1卷积层和四个不同空洞率的卷积层以及一个池化层。这些层的输出在通道维度上拼接起来，然后通过一个1x1卷积层进行融合，最终输出特征图。

modules = []
modules.append(nn.Sequential(
nn.Conv2d(in_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU()))
这一段表示定义了一个数组modules，并往里面塞了一个1*1卷积，一个批归一化和一个ReLU函数，对应结构图中的1*1卷积

rate1, rate2, rate3 = tuple(atrous_rates)
modules.append(ASPPConv(in_channels, out_channels, rate1))
modules.append(ASPPConv(in_channels, out_channels, rate2))
modules.append(ASPPConv(in_channels, out_channels, rate3))
modules.append(ASPPPooling(in_channels, out_channels))

这段代码表示往modules中加入四个值，分别是rate1，rate2，rate3空洞率下的空洞卷积，以及池化后的图像。

self.convs = nn.ModuleList(modules)

nn.ModuleList 是一个容器，可以将多个 nn.Module 对象组合在一起，方便进行参数管理和优化器更新等操作。

在这个代码片段中，self.convs 是一个 nn.ModuleList 对象，它包含了若干个 nn.Module 对象，这些 nn.Module 对象可能是卷积层、池化层、全连接层等等。这个 nn.ModuleList 对象可以被用于搭建神经网络模型。

self.project = nn.Sequential(
nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
nn.Dropout(0.5))

这段代码是一个卷积神经网络的一部分，其中包含了一个卷积层、一个批量归一化层、一个ReLU激活函数和一个Dropout层。这个网络的输入是一个5倍于输出通道数的张量，经过卷积层后输出通道数为out_channels。

主要作用是对ASPP五个尺度融合后的通道数调整

在前向传播过程中

def forward(self, x):
res = []
for conv in self.convs:
res.append(conv(x))
res = torch.cat(res, dim=1)
return self.project(res)

定义一个数组res，继承self.convs中的值，就是将ASPP中的五个尺度进行concatenate融合，之后再进行1*1卷积通道调整。

完整代码

from torch import nn
import torch
import torch.nn.functional as F
 
class ASPPConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, dilation):
        modules = [
            nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        ]
        super(ASPPConv, self).__init__(*modules)
 
class ASPPPooling(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(ASPPPooling, self).__init__(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU())
 
    def forward(self, x):
        size = x.shape[-2:]
        x = super(ASPPPooling, self).forward(x)
        return F.interpolate(x, size=size, mode='bilinear', align_corners=False)
 
class ASPP(nn.Module):
    def __init__(self, in_channels, atrous_rates):
        super(ASPP, self).__init__()
        out_channels = 256
        modules = []
        modules.append(nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()))
 
        rate1, rate2, rate3 = tuple(atrous_rates)
        modules.append(ASPPConv(in_channels, out_channels, rate1))
        modules.append(ASPPConv(in_channels, out_channels, rate2))
        modules.append(ASPPConv(in_channels, out_channels, rate3))
        modules.append(ASPPPooling(in_channels, out_channels))
 
        self.convs = nn.ModuleList(modules)
 
        self.project = nn.Sequential(
            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(0.5))
 
    def forward(self, x):
        res = []
        for conv in self.convs:
            res.append(conv(x))
        res = torch.cat(res, dim=1)
        return self.project(res)
 
aspp = ASPP(256,[6,12,18])
x = torch.rand(2,256,13,13)
print(aspp(x).shape)

乐奇 Rokid 开放社区

更多推荐

深蹲做到第 7 下，眼镜该显示什么？《铁面私教》的 Rokid AIUI 开发实战

大家好，我是 Steven。深蹲做到第七下，你想知道还差几下。等把杠铃放回架子上，你才有空看刚才这一组练得怎么样、还要休息多久。两次抬眼，中间可能只隔了几秒，需要的却是两种界面。我做「铁面私教」，就是想把这两次抬眼接起来。它是一款面向独自做力量训练的人的眼镜陪练。训练中看次数，组间看点评，练完把记录存进日志。铁面私教 · Rokid Glasses 原生 AIUI 智能体 · 商店已上架 v0.3