好的，下面是一个使用 PyTorch 实现 UNet-MobileNet 的代码示例。我们将 MobileNet 作为编码器（Encoder），并构建一个标准的 UNet 结构。 ### UNet-MobileNet 实现 ```python import torch import torch.nn as nn from torchvision import models class UNetMobileNet(nn.Module): def __init__(self, num_classes=1, pretrained=True): super(UNetMobileNet, self).__init__() # 使用 MobileNetV2 作为编码器 mobilenet = models.mobilenet_v2(pretrained=pretrained) self.encoder = mobilenet.features # 获取 MobileNet 的各个阶段输出通道数 filters = [16, 24, 32, 64, 1280] # MobileNetV2 的特征图通道数 # 解码器部分 self.decoder_blocks = nn.ModuleList([ self._decoder_block(filters[-1], filters[-2]), self._decoder_block(filters[-2], filters[-3]), self._decoder_block(filters[-3], filters[-4]), self._decoder_block(filters[-4], filters[-5]), ]) # 最终卷积层将通道数转换为目标类别数 self.final_conv = nn.Conv2d(filters[0], num_classes, kernel_size=1) def _decoder_block(self, in_channels, out_channels): """定义解码器块""" return nn.Sequential( nn.ConvTranspose2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): # 编码器部分：提取特征 skips = [] for i in range(len(self.encoder)): x = self.encoder[i](x) if i in [2, 4, 7, 14]: # 保存中间层的特征图用于跳跃连接 skips.append(x) # 最后一层特征图 x = self.encoder[-1](x) # 解码器部分：上采样并融合特征 for i, decoder_block in enumerate(self.decoder_blocks): x = decoder_block(x) # 跳跃连接 if i < len(skips): x = torch.cat([x, skips[-(i+1)]], dim=1) # 最终卷积层 x = self.final_conv(x) return x # 测试模型 if __name__ == "__main__": model = UNetMobileNet(num_classes=1) print(model) # 测试输入和输出 input_tensor = torch.randn(1, 3, 224, 224) # 假设输入图像大小为 224x224 output = model(input_tensor) print("Output shape:", output.shape) ``` --- ### 上述代码解释 1. **编码器部分**: - 我们使用 `torchvision.models.mobilenet_v2` 提供的预训练 MobileNetV2 模型作为编码器。 - MobileNetV2 是一种轻量级网络，适合用作语义分割任务中的编码器。 - 我们从 MobileNetV2 的不同阶段提取特征图，并将其存储在 `skips` 列表中以备后续跳跃连接使用。 2. **解码器部分**: - 解码器由多个反卷积层（`ConvTranspose2d`）组成，用于逐步上采样特征图。 - 每个解码器块通过跳跃连接与编码器的对应特征图融合，从而保留更多的细节信息。 3. **最终卷积层**: - 最后的卷积层将通道数调整为目标类别的数量（例如二分类任务时 `num_classes=1`）。 4. **跳跃连接**: - 在解码器中，我们通过跳跃连接将编码器的中间特征图与当前解码器的特征图拼接在一起，以增强特征表达能力。 5. **测试模型**: - 我们创建了一个简单的测试用例，输入张量形状为 `(1, 3, 224, 224)`，输出形状应为 `(1, 1, 224, 224)`。 --- ### 注意事项 1. **跳跃连接的层数**: - 在代码中，我们根据 MobileNetV2 的结构选择了一些关键层进行跳跃连接。可以根据具体任务需求调整这些层的选择。 2. **输出尺寸**: - 输出特征图的尺寸应与输入图像的尺寸一致。如果需要更大的输入图像，可以通过调整解码器的上采样参数来实现。 3. **预训练权重**: - 如果设置 `pretrained=True`，则会加载 ImageNet 预训练权重，这有助于提高模型的初始性能。 --- ###