解决PyTorch多任务模型中批次大小不一致问题：卷积层输出展平与全连接层连接

多任务分类模型构建挑战

在深度学习领域，有时我们需要一个模型同时完成多个相关的分类任务，例如，给定一幅图像，同时预测其艺术家、流派和风格。这被称为多任务分类。构建此类模型时，通常有两种策略：

1. 修改预训练模型： 利用像Hugging Face Transformers库中提供的预训练模型（如ResNet18），替换或添加自定义的分类头。这种方法通常需要理解预训练模型的内部结构，以确保新添加的层能正确连接到模型的特征提取部分。
2. 构建自定义模型： 从零开始或基于简单的骨干网络构建一个全新的模型，其中包含共享的特征提取层和针对每个任务的独立分类分支。

在实践中，直接修改预训练模型（如ResNet18）的分类器可能不如预期。例如，简单地为ResNetForImageClassification实例添加classifier_artist、classifier_style、classifier_genre等属性，并不能自动将其集成到模型的forward方法中。torchinfo的输出也印证了这一点，模型的主体仍然是其原有的ResNetModel和Sequential (classifier)，并未包含新定义的分类器。这通常意味着需要继承并重写模型的forward方法，或者正确地替换原有的分类头。

当自定义PyTorch模型时，我们拥有更大的灵活性来设计多任务架构。然而，这也引入了新的挑战，尤其是在处理不同层之间的数据维度匹配问题上。

批次大小不一致问题分析

构建自定义的WikiartModel用于多任务分类时，我们定义了共享的卷积层用于特征提取，并为艺术家、流派和风格三个任务分别设置了独立的全连接层分支。模型定义如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

class WikiartModel(nn.Module):
    def __init__(self, num_artists, num_genres, num_styles):
        super(WikiartModel, self).__init__()

        # Shared Convolutional Layers
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

        # Artist classification branch (Incorrect input size)
        self.fc_artist1 = nn.Linear(256 * 16 * 16, 512) # Potentially incorrect
        self.fc_artist2 = nn.Linear(512, num_artists)

        # Genre classification branch (Incorrect input size)
        self.fc_genre1 = nn.Linear(256 * 16 * 16, 512) # Potentially incorrect
        self.fc_genre2 = nn.Linear(512, num_genres)

        # Style classification branch (Incorrect input size)
        self.fc_style1 = nn.Linear(256 * 16 * 16, 512) # Potentially incorrect
        self.fc_style2 = nn.Linear(512, num_styles)

    def forward(self, x):
        # Shared convolutional layers
        x = self.pool(F.relu(self.conv1(x)))   
        x = self.pool(F.relu(self.conv2(x)))       
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 256 * 16 * 16) # Potentially incorrect flattening

        # Artist classification branch
        artists_out = F.relu(self.fc_artist1(x))
        artists_out = self.fc_artist2(artists_out)

        # Genre classification branch
        genre_out = F.relu(self.fc_genre1(x))
        genre_out = self.fc_genre2(genre_out) 

        # Style classification branch 
        style_out = F.relu(self.fc_style1(x))
        style_out = self.fc_style2(style_out)

        return artists_out, genre_out, style_out

# num_artists, num_genres, num_styles are defined externally

在使用torchinfo检查模型结构时，我们发现一个关键问题：模型的输入批次大小为32（例如[32, 3, 224, 224]），但其内部全连接层（如fc_artist1）的输入批次大小却变成了98，导致最终输出的批次大小也为98。这直接引发了训练循环中计算损失时的ValueError: Expected input batch_size (98) to match target batch_size (32).错误。

问题根源分析：

这个批次大小不一致的根本原因在于卷积层输出特征图的尺寸计算错误，以及随后对特征图进行展平（flatten）操作时，全连接层期望的输入维度与实际不符。

让我们逐步分析数据流：

1. 初始输入： [Batch_Size, 3, 224, 224] (假设 Batch_Size = 32)
2. self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1):
   • 输入：[32, 3, 224, 224]
   • 输出：[32, 64, 224, 224] (由于 padding=1, 尺寸不变)
3. x = self.pool(F.relu(self.conv1(x))) (self.pool = nn.MaxPool2d(2, 2)):
   • 输入：[32, 64, 224, 224]
   • 输出：[32, 64, 112, 112] (尺寸减半)
4. self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1):
   • 输入：[32, 64, 112, 112]
   • 输出：[32, 128, 112, 112]
5. x = self.pool(F.relu(self.conv2(x))):
   • 输入：[32, 128, 112, 112]
   • 输出：[32, 128, 56, 56]
6. self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1):
   • 输入：[32, 128, 56, 56]
   • 输出：[32, 256, 56, 56]
7. x = self.pool(F.relu(self.conv3(x))):
   • 输入：[32, 256, 56, 56]
   • 输出：[32, 256, 28, 28]

因此，在进入全连接层之前，特征图的实际尺寸是 [32, 256, 28, 28]。

问题出在这一行：x = x.view(-1, 256 * 16 * 16)。 当x的实际形状是[32, 256, 28, 28]时，总元素数量为 32 * 256 * 28 * 28 = 6422528。 而256 * 16 * 16 = 65536。 当使用x.view(-1, 65536)时，PyTorch会尝试将总元素数量除以65536来推断-1对应的维度： 6422528 / 65536 = 98。 所以，x被错误地展平为了[98, 65536]，导致批次大小从32变成了98。

解决方案：正确计算与展平特征图

要解决这个问题，我们需要确保全连接层的输入维度与卷积层输出的实际展平尺寸相匹配，并且批次大小在展平过程中保持不变。

火山方舟

火山引擎一站式大模型服务平台，已接入满血版DeepSeek

下载

步骤一：确定卷积层最终输出尺寸

如上分析，经过三次卷积和三次最大池化操作后，对于 224x224 的输入图像，最终的特征图尺寸是 [Batch_Size, 256, 28, 28]。因此，展平后的特征向量长度应该是 256 * 28 * 28。

步骤二：正确展平操作

在将卷积层的输出传递给全连接层之前，需要将其展平为二维张量 [Batch_Size, Features]。为了确保批次大小不变，应该使用 x.view(x.size(0), -1)。这里的 x.size(0) 会保留原始的批次大小（例如32），而 -1 会自动计算剩余维度的乘积，将其展平为单个特征向量。

对于 [32, 256, 28, 28] 的张量，x.view(x.size(0), -1) 会将其展平为 [32, 256 * 28 * 28]，即 [32, 200704]。

步骤三：修正全连接层输入维度

基于正确的展平尺寸，所有连接到卷积层输出的全连接层（fc_artist1, fc_genre1, fc_style1）的 in_features 参数都应该修改为 256 * 28 * 28。

# 将 nn.Linear(256 * 16 * 16, 512)
# 修正为
nn.Linear(256 * 28 * 28, 512) # 256 * 28 * 28 = 200704

修正后的WikiartModel代码示例

根据上述修正，WikiartModel的定义应更新如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

class WikiartModel(nn.Module):
    def __init__(self, num_artists, num_genres, num_styles):
        super(WikiartModel, self).__init__()

        # Shared Convolutional Layers
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

        # 计算卷积层最终输出的特征图尺寸，用于全连接层
        # 对于224x224输入，经过三次conv+pool后，尺寸变为 28x28
        self.final_feature_map_size = 28 
        self.flattened_features = 256 * self.final_feature_map_size * self.final_feature_map_size # 256 * 28 * 28 = 200704

        # Artist classification branch
        self.fc_artist1 = nn.Linear(self.flattened_features, 512)
        self.fc_artist2 = nn.Linear(512, num_artists)

        # Genre classification branch
        self.fc_genre1 = nn.Linear(self.flattened_features, 512)
        self.fc_genre2 = nn.Linear(512, num_genres)

        # Style classification branch
        self.fc_style1 = nn.Linear(self.flattened_features, 512) 
        self.fc_style2 = nn.Linear(512, num_styles)

    def forward(self, x):
        # Shared convolutional layers
        x = self.pool(F.relu(self.conv1(x)))   # Output: [Batch_Size, 64, 112, 112]
        x = self.pool(F.relu(self.conv2(x)))   # Output: [Batch_Size, 128, 56, 56]
        x = self.pool(F.relu(self.conv3(x)))   # Output: [Batch_Size, 256, 28, 28]

        # Correct flattening: preserve batch size, flatten remaining dimensions
        x = x.view(x.size(0), -1) # Output: [Batch_Size, 256 * 28 * 28] = [Batch_Size, 200704]

        # Artist classification branch
        artists_out = F.relu(self.fc_artist1(x))
        artists_out = self.fc_artist2(artists_out)

        # Genre classification branch
        genre_out = F.relu(self.fc_genre1(x))
        genre_out = self.fc_genre2(genre_out) 

        # Style classification branch 
        style_out = F.relu(self.fc_style1(x))
        style_out = self.fc_style2(style_out)

        return artists_out, genre_out, style_out

# Example usage:
num_artists = 129
num_genres = 11
num_styles = 27

model = WikiartModel(num_artists, num_genres, num_styles)
# Now, if you pass a tensor of shape [32, 3, 224, 224] to the model,
# the outputs will correctly have a batch size of 32.
# e.g., artists_out.shape will be [32, 129]

总结与注意事项

批次大小不一致是PyTorch模型开发中常见的错误，尤其是在卷积层和全连接层之间进行维度转换时。解决此问题的关键在于：

• 精确计算中间层输出尺寸： 在设计网络时，务必仔细推导每个卷积层和池化层的输出尺寸。对于图像数据，常用的计算公式为 (输入尺寸 - 卷积核尺寸 + 2 * 填充) / 步长 + 1。
• 正确使用展平操作： 当需要将多维特征图展平为一维向量以供全连接层使用时，始终推荐使用 tensor.view(tensor.size(0), -1)。这能确保批次维度保持不变，而其余维度则被正确地展平。
• 匹配全连接层输入维度： 全连接层（nn.Linear）的 in_features 参数必须与前一层输出的展平特征向量的长度完全匹配。
• 利用调试工具： 在模型构建和调试阶段，积极使用 torchinfo.summary() 或在 forward 方法中打印 tensor.shape，能够直观地检查每一层的数据流和尺寸变化，从而快速定位维度不匹配问题。

通过遵循这些原则，可以有效地避免和解决PyTorch模型中因维度不匹配导致的批次大小不一致问题，确保模型能够顺利训练。