efficientnet迁移至语义分割的pytorch代码，附中文注释

以下是 efficientnet 迁移至语义分割的 PyTorch 代码，带有中文注释：

import torch
import torch.nn as nn
from efficientnet_pytorch import EfficientNet

class ConvBlock(nn.Module):
    # 定义卷积块
    def __init__(self, in_ch, out_ch):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        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 DecoderBlock(nn.Module):
    # 定义解码块
    def __init__(self, in_ch, out_ch):
        super(DecoderBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_ch, in_ch // 4, kernel_size=1, stride=1)
        self.bn1 = nn.BatchNorm2d(in_ch // 4)
        self.relu1 = nn.ReLU(inplace=True)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv2 = ConvBlock(in_ch // 4, out_ch)

    def forward(self, x, e):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.upsample(x)
        x = torch.cat([x, e], dim=1)
        x = self.conv2(x)
        return x

class EfficientUNet(nn.Module):
    # 定义 EfficientUNet 模型
    def __init__(self, num_classes):
        super(EfficientUNet, self).__init__()
        self.encoder = EfficientNet.from_pretrained('efficientnet-b0')
        self.conv1 = ConvBlock(320, 256)
        self.conv2 = ConvBlock(192, 128)
        self.conv3 = ConvBlock(112, 64)
        self.conv4 = ConvBlock(40, 32)
        self.center = DecoderBlock(1280, 512)
        self.dec5 = DecoderBlock(512 + 320, 256)
        self.dec6 = DecoderBlock(256 + 192, 128)
        self.dec7 = DecoderBlock(128 + 112, 64)
        self.dec8 = DecoderBlock(64 + 40, 32)
        self.final_conv = nn.Conv2d(32, num_classes, kernel_size=1)

    def forward(self, x):
        # 编码器部分
        x1 = self.encoder.extract_features(x)
        x2 = self.encoder._swish(self.encoder._bn1(self.encoder._conv_head(x1)))
        x3 = self.encoder._swish(self.encoder._bn2(self.encoder._conv_stem(x2)))
        x4 = self.encoder._swish(self.encoder._bn3(self.encoder._blocks[0](x3)))
        x5 = self.encoder._swish(self.encoder._bn4(self.encoder._blocks[1](x4)))
        x6 = self.encoder._swish(self.encoder._bn5(self.encoder._blocks[2](x5)))
        x7 = self.encoder._swish(self.encoder._bn6(self.encoder._blocks[3](x6)))
        x8 = self.encoder._swish(self.encoder._bn7(self.encoder._blocks[4](x7)))
        x9 = self.encoder._swish(self.encoder._bn8(self.encoder._conv_head(x8)))
        # 解码器部分
        x = self.center(x9)
        x = self.dec5(x, x8)
        x = self.dec6(x, x7)
        x = self.dec7(x, x6)
        x = self.dec8(x, x4)
        out = self.final_conv(x)
        return out

此代码将 EfficientNet 模型迁移至语义分割任务中，使用了 U-Net 结构。其中 `ConvBlock` 和 `DecoderBlock` 分别定义了卷积块和解码块，`EfficientUNet` 则定义了整个网络结构。在 `EfficientUNet` 中，首先使用 EfficientNet 作为编码器，然后使用 `DecoderBlock` 将编码器输出的特征图进行解码。最后使用一个卷积层将解码后的特征图输出为语义分割结果。