如何为segmentation-models-pytorch的编码器添加自定义分类输出头

要给基于segmentation-models-pytorch的MAnet模型编码器添加自定义分类输出，你可以通过构建多任务模型的方式，复用原模型的编码器，再额外添加分类头实现。具体步骤如下：

1. 导入依赖库 #

import torch
import torch.nn as nn
import segmentation_models_pytorch as smp

2. 定义多任务模型 #

自定义一个继承自nn.Module的模型类，整合原MAnet的分割模块和新增的分类头：

class MultiTaskModel(nn.Module):
    def __init__(self, encoder_name="efficientnet-b0", encoder_weights="imagenet", in_channels=3, seg_classes=1, cls_classes=10):
        super().__init__()
        # 初始化MAnet分割模型
        self.segmentation_model = smp.MAnet(
            encoder_name=encoder_name,
            encoder_weights=encoder_weights,
            in_channels=in_channels,
            classes=seg_classes,
        )
        # 获取编码器最后一层特征的维度（不同编码器对应值不同）
        encoder_final_dim = self.segmentation_model.encoder.out_channels[-1]
        # 构建自定义分类头
        self.classification_head = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),  # 全局平均池化
            nn.Flatten(),  # 展平特征
            nn.Linear(encoder_final_dim, cls_classes),  # 全连接层输出分类结果
            # 根据任务类型选择激活函数：多分类用nn.Softmax(dim=1)，二分类用nn.Sigmoid()
        )
        
    def forward(self, x):
        # 计算分割输出
        seg_output = self.segmentation_model(x)
        # 获取编码器最后一层的特征图
        encoder_last_feat = self.segmentation_model.encoder(x)[-1]
        # 计算分类输出
        cls_output = self.classification_head(encoder_last_feat)
        # 返回分割和分类两个结果
        return seg_output, cls_output

3. 使用模型示例 #

# 初始化模型，这里假设分类任务需要输出10类
model = MultiTaskModel(cls_classes=10)
# 构造测试输入（batch_size=2, 3通道, 256x256尺寸）
test_input = torch.randn(2, 3, 256, 256)
# 前向传播得到输出
seg_out, cls_out = model(test_input)

print(f"分割输出形状: {seg_out.shape}")  # 应为(2, 1, 256, 256)
print(f"分类输出形状: {cls_out.shape}")  # 应为(2, 10)

注意事项 #

• 不同编码器的最后一层特征维度不同，可以通过print(self.segmentation_model.encoder.out_channels)查看具体数值，比如efficientnet-b0对应1280，resnet50对应2048。
• 分类头的结构可根据需求调整，比如添加nn.Dropout(p=0.5)防止过拟合，或者增加额外的全连接层。
• 训练时需要同时定义分割损失（如DiceLoss、CrossEntropyLoss）和分类损失，将两者加权求和作为总损失进行反向传播。

内容的提问来源于stack exchange，提问作者Anton Ganichev