如何修改Densenet121模型的out_features？代码运行报错求助

1. 报错原因分析 #

报错提示expected input[64, 64, 62, 62] to have 3 channels, but got 64 channels instead，核心问题有两个：

• 预训练Densenet121的第一层卷积层默认只接受3通道输入，但你自定义的cnn1输出了64通道特征图，直接传入导致通道不匹配。
• 你在forward方法中调用densenet()动态修改模型结构，这会导致重复初始化、参数冻结操作，破坏前向传播的稳定性。

2. 正确修改方案 #

步骤1：在初始化阶段完成模型配置 #

在__init__中一次性完成Densenet的参数冻结、输入通道适配、输出维度修改，避免在forward中动态调整模型：

• 冻结预训练参数：直接遍历模型参数设置requires_grad=False
• 适配输入通道：替换Densenet的第一层卷积，将输入通道改为64（匹配cnn1的输出）
• 修改输出维度：替换原分类器层，设置目标out_features

步骤2：修正forward流程的维度匹配问题 #

• 移除forward中错误的densenet()调用，直接使用初始化好的Densenet特征提取部分
• 修正批归一化层类型（全连接层后用BatchNorm1d而非BatchNorm2d）
• 调整特征展平的时机，确保全连接层输入维度正确

修改后的完整代码 #

import torch
import torch.nn as nn
from torchvision import models
from torch.nn import AvgPool2d

class Densnet121(nn.Module):
    def __init__(self):
        super(Densnet121, self).__init__() 
        self.cnn1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1)
        # 初始化预训练Densenet121
        self.Densenet_121 = models.densenet121(pretrained=True)
        
        # 冻结预训练参数
        for param in self.Densenet_121.parameters():
            param.requires_grad = False
        
        # 修改Densenet输入通道（匹配cnn1的64通道输出）
        original_conv0 = self.Densenet_121.features.conv0
        self.Densenet_121.features.conv0 = nn.Conv2d(
            in_channels=64, 
            out_channels=original_conv0.out_channels,
            kernel_size=original_conv0.kernel_size,
            stride=original_conv0.stride,
            padding=original_conv0.padding,
            bias=original_conv0.bias
        )
        
        # 修改Densenet输出特征维度（这里设置为1024，可按需调整为目标out_features）
        self.Densenet_121.classifier = nn.Linear(1024, 1024)
        
        # 后续层适配维度
        self.gap = AvgPool2d(kernel_size=2, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(1024)
        self.do1 = nn.Dropout(0.25)
        self.linear = nn.Linear(1024, 256)
        self.bn2 = nn.BatchNorm1d(256)  # 全连接后用1D批归一化
        self.do2 = nn.Dropout(0.25)
        self.output = nn.Linear(256, 2)
        self.act = nn.ReLU()
        
    def forward(self, x):
        # 自定义卷积层
        img = self.act(self.cnn1(x))
        # 传入修改后的Densenet特征提取部分
        img = self.Densenet_121.features(img)
        
        # 特征处理流程
        img = self.gap(img)
        img = self.bn1(img)
        img = self.do1(img)
        img = torch.flatten(img, 1)  # 展平为一维特征
        
        img = self.linear(img)
        img = self.act(img)
        img = self.bn2(img)
        img = self.do2(img)
        img = self.output(img)
        
        return img

3. 关键修改点说明 #

• 输入通道适配：替换Densenet的conv0层，彻底解决自定义卷积与预训练模型的通道不匹配问题。
• 输出维度修改：直接替换classifier层为nn.Linear(1024, 目标维度)，即可灵活设置out_features。
• 层类型修正：将bn2改为BatchNorm1d，适配全连接层输出的一维特征格式。
• 流程优化：所有模型结构调整放在__init__中，保证forward仅执行前向计算，避免重复操作导致的错误。

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