aspp unet代码

以下是 ASP-UNet 的 PyTorch 代码：

import torch
import torch.nn as nn

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

class ASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ASPP, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=6, dilation=6)
        self.conv3 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=12, dilation=12)
        self.conv4 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=18, dilation=18)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv5 = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x3 = self.conv3(x)
        x4 = self.conv4(x)
        x5 = self.avg_pool(x)
        x5 = self.conv5(x5)
        x5 = nn.functional.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat([x1, x2, x3, x4, x5], dim=1)
        x = self.relu(x)
        return x

class Decoder(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Decoder, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels//4, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels//4, out_channels, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x, skip):
        x = nn.functional.interpolate(x, size=skip.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat([x, skip], dim=1)
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        return x

class ASP_UNet(nn.Module):
    def __init__(self, num_classes):
        super(ASP_UNet, self).__init__()
        self.first_conv = ConvBlock(3, 64)
        self.encoder1 = ConvBlock(64, 128)
        self.encoder2 = ConvBlock(128, 256)
        self.encoder3 = ConvBlock(256, 512)
        self.encoder4 = ConvBlock(512, 1024)
        self.aspp = ASPP(1024, 256)
        self.decoder1 = Decoder(256+1024, 512)
        self.decoder2 = Decoder(128+512, 256)
        self.decoder3 = Decoder(64+256, 128)
        self.decoder4 = Decoder(128, 64)
        self.last_conv = nn.Conv2d(64, num_classes, kernel_size=1)

    def forward(self, x):
        x1 = self.first_conv(x)
        x2 = self.encoder1(x1)
        x3 = self.encoder2(x2)
        x4 = self.encoder3(x3)
        x = self.encoder4(x4)
        x = self.aspp(x)
        x = self.decoder1(x, x4)
        x = self.decoder2(x, x3)
        x = self.decoder3(x, x2)
        x = self.decoder4(x, x1)
        x = self.last_conv(x)
        x = nn.functional.interpolate(x, size=x.size()[2:], mode='bilinear', align_corners=True)
        return x

这是基于 PyTorch 实现的 ASP-UNet 模型代码，包括 ConvBlock、ASPP、Decoder 和 ASP_UNet 等类。其中 ASP_UNet 是主模型类，包括了整个网络的前向传播过程。您可以根据自己的需要进行修改和使用。