27、ASPP模塊

論文《Rethinking Atrous Convolution for Semantic Image Segmentation》

1、作用

DeepLabv3是一種先進(jìn)的語(yǔ)義圖像分割系統(tǒng)，它通過(guò)使用空洞卷積捕獲多尺度上下文來(lái)顯著提升性能，無(wú)需依賴(lài)DenseCRF后處理。

2、機(jī)制

DeepLabv3的核心機(jī)制圍繞空洞（擴(kuò)張）卷積展開(kāi)。這種技術(shù)允許模型控制濾波器的視野，使其能夠在多個(gè)尺度上捕獲空間上下文。DeepLabv3在串聯(lián)和并聯(lián)架構(gòu)中使用空洞卷積來(lái)提取密集的特征圖，并有效地整合多尺度信息。文章還介紹了Atrous Spatial Pyramid Pooling（ASPP）模塊，該模塊通過(guò)在多個(gè)尺度上探索卷積特征并結(jié)合圖像級(jí)特征，用于編碼全局上下文。

3、獨(dú)特優(yōu)勢(shì)

1、多尺度上下文捕獲：

通過(guò)在不同配置中使用空洞卷積，DeepLabv3能夠從多個(gè)尺度捕獲上下文信息，這對(duì)于準(zhǔn)確分割不同大小的對(duì)象至關(guān)重要。

2、高效密集特征提取：

空洞卷積使得模型能夠在不需要額外參數(shù)或計(jì)算資源的情況下提取密集特征圖，提高了部署效率。

3、性能提升：

ASPP與圖像級(jí)特征的結(jié)合顯著提高了模型性能，使其在PASCAL VOC 2012等基準(zhǔn)數(shù)據(jù)集上與其他最先進(jìn)方法競(jìng)爭(zhēng)。

4、靈活性和泛化能力：

DeepLabv3的框架是通用的，可以應(yīng)用于任何網(wǎng)絡(luò)架構(gòu)，為適應(yīng)不同的分割任務(wù)提供了靈活性。

4、代碼

from torch import nn
import torch
import torch.nn.functional as F


# 定義一個(gè)包含空洞卷積、批量歸一化和ReLU激活函數(shù)的子模塊
class ASPPConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, dilation):
        modules = [
            # 空洞卷積，通過(guò)調(diào)整dilation參數(shù)來(lái)捕獲不同尺度的信息
            nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
            nn.BatchNorm2d(out_channels),  # 批量歸一化
            nn.ReLU()  # ReLU激活函數(shù)
        ]
        super(ASPPConv, self).__init__(*modules)


# 定義一個(gè)全局平均池化后接卷積、批量歸一化和ReLU的子模塊
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),  # 1x1卷積
            nn.BatchNorm2d(out_channels),  # 批量歸一化
            nn.ReLU())  # ReLU激活函數(shù)

    def forward(self, x):
        size = x.shape[-2:]  # 保存輸入特征圖的空間維度
        x = super(ASPPPooling, self).forward(x)
        # 通過(guò)雙線(xiàn)性插值將特征圖大小調(diào)整回原始輸入大小
        return F.interpolate(x, size=size, mode='bilinear', align_corners=False)


# ASPP模塊主體，結(jié)合不同膨脹率的空洞卷積和全局平均池化
class ASPP(nn.Module):
    def __init__(self, in_channels, atrous_rates):
        super(ASPP, self).__init__()
        out_channels = 256  # 輸出通道數(shù)
        modules = []
        modules.append(nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),  # 1x1卷積用于降維
            nn.BatchNorm2d(out_channels),
            nn.ReLU()))

        # 根據(jù)不同的膨脹率添加空洞卷積模塊
        for rate in atrous_rates:
            modules.append(ASPPConv(in_channels, out_channels, rate))

        # 添加全局平均池化模塊
        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))  # 防止過(guò)擬合的Dropout層

    def forward(self, x):
        res = []
        # 對(duì)每個(gè)模塊的輸出進(jìn)行收集
        for conv in self.convs:
            res.append(conv(x))
        # 將收集到的特征在通道維度上拼接
        res = torch.cat(res, dim=1)
        # 對(duì)拼接后的特征進(jìn)行處理
        return self.project(res)


# 示例使用ASPP模塊
aspp = ASPP(256, [6, 12, 18])
x = torch.rand(2, 256, 13, 13)
print(aspp(x).shape)  # 輸出處理后的特征圖維度