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)