【Bug已解决】PyTorch model input shape 解决方案问题描述在 PyTorch 中构建和训练神经网络时正确理解和处理模型的输入形状input shape是最基础但也最容易出错的问题之一。许多开发者在定义模型、准备数据、进行前向传播时经常因为输入形状不匹配而遇到各种错误。典型的问题场景包括RuntimeError: mat1 and mat2 shapes cannot be multiplied—— 全连接层输入维度不匹配RuntimeError: Expected 3-dimensional (unbatched) or 4-dimensional (batched) input—— 卷积层输入维度错误RuntimeError: Given groups1, weight of size ... expected input ...—— 卷积核通道数不匹配DataLoader 输出的 batch 形状与模型期望不一致单个样本和批量样本的形状混淆在forward方法中形状变换错误导致后续层报错这些问题的核心在于理解 PyTorch 中 tensor 的维度约定以及不同层对输入形状的要求。错误复现场景一全连接层维度不匹配import torch import torch.nn as nn model nn.Linear(784, 10) # 期望输入最后一维是 784 x torch.randn(32, 100) # 但输入是 100 output model(x) # RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x100 and 784x10)场景二卷积层输入维度错误conv nn.Conv2d(3, 16, kernel_size3, padding1) # 期望 4D 输入 [B, C, H, W] x torch.randn(3, 32, 32) # 但输入是 3D缺少 batch 维度 output conv(x) # RuntimeError: Expected 3-dimensional (unbatched) or 4-dimensional (batched) input场景三卷积通道数不匹配conv nn.Conv2d(3, 16, kernel_size3) # 期望输入通道数为 3 x torch.randn(4, 1, 32, 32) # 但输入通道数为 1 output conv(x) # RuntimeError: Given groups1, weight of size [16, 3, 3, 3], # expected input with 3 channels, but got 1 channels场景四RNN 输入形状错误lstm nn.LSTM(input_size10, hidden_size20, batch_firstTrue) # 期望输入: [batch_size, seq_len, input_size] x torch.randn(10, 100, 5) # seq_len100, input_size5不是10 output, _ lstm(x) # RuntimeError: input.size(-1) must be equal to input_size场景五batch 维度混淆model nn.Sequential( nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10), ) # 单个样本忘记添加 batch 维度 x torch.randn(784) # 1D tensor output model(x) # 可能不报错但结果错误因为 model 把 784 当作 batch_size # 正确添加 batch 维度 x torch.randn(1, 784) # 2D tensor output model(x)根因分析1. PyTorch 的维度约定PyTorch 中不同类型的层对输入形状有不同的约定层类型期望输入形状说明nn.Linear[batch_size, in_features]2Dnn.Conv1d[batch_size, in_channels, seq_len]3Dnn.Conv2d[batch_size, in_channels, height, width]4Dnn.Conv3d[batch_size, in_channels, depth, height, width]5Dnn.RNN/GRU/LSTM[seq_len, batch_size, input_size]或[batch_size, seq_len, input_size]batch_firstTrue3Dnn.BatchNorm1d[batch_size, num_features]或[batch_size, num_features, seq_len]2D/3Dnn.BatchNorm2d[batch_size, num_channels, height, width]4Dnn.Embedding[batch_size]或[batch_size, seq_len]整数索引2. batch 维度的重要性PyTorch 的几乎所有层都期望输入的第一个维度是 batch_size。即使是单个样本也需要添加一个大小为 1 的 batch 维度# 单个样本 x_single torch.randn(3, 32, 32) # [C, H, W] x_batched x_single.unsqueeze(0) # [1, C, H, W]3. 形状变换的链式效应在nn.Sequential或forward方法中前一层的输出形状必须与后一层的期望输入形状匹配。一个环节的错误会传播到后续所有层。4. DataLoader 的形状影响DataLoader会在第 0 维添加 batch 维度。如果Dataset.__getitem__返回形状[C, H, W]的 tensorDataLoader输出形状为[B, C, H, W]。解决方案方案一使用形状检查工具import torch import torch.nn as nn from typing import Tuple, Optional class ShapeChecker: 模型输入/输出形状检查器 staticmethod def check_linear_input(model, input_shape): 检查 Linear 层输入 if isinstance(model, nn.Linear): expected model.in_features actual input_shape[-1] if actual ! expected: raise ValueError( fLinear layer expects {expected} features, fbut got {actual} ) staticmethod def check_conv2d_input(model, input_shape): 检查 Conv2d 层输入 if isinstance(model, nn.Conv2d): expected model.in_channels if len(input_shape) 4: actual input_shape[1] elif len(input_shape) 3: actual input_shape[0] else: raise ValueError(fConv2d expects 3D or 4D input, got {len(input_shape)}D) if actual ! expected: raise ValueError( fConv2d expects {expected} channels, but got {actual} ) staticmethod def trace_shapes(model, input_tensor, layer_namesNone): 追踪模型每层的输入输出形状 shapes [] hooks [] def hook_fn(module, input, output): input_shape tuple(input[0].shape) output_shape tuple(output.shape) shapes.append({ layer: module.__class__.__name__, input_shape: input_shape, output_shape: output_shape, }) # 注册 hook for module in model.modules(): if len(list(module.children())) 0: # 叶子模块 h module.register_forward_hook(hook_fn) hooks.append(h) # 前向传播 with torch.no_grad(): _ model(input_tensor) # 移除 hook for h in hooks: h.remove() return shapes staticmethod def print_model_summary(model, input_tensor): 打印模型摘要类似 Keras summary shapes ShapeChecker.trace_shapes(model, input_tensor) print(f{Layer:25} {Input Shape:25} {Output Shape:25}) print(- * 75) total_params 0 for info in shapes: layer_name info[layer] in_shape str(info[input_shape]) out_shape str(info[output_shape]) print(f{layer_name:25} {in_shape:25} {out_shape:25}) print(- * 75) total_params sum(p.numel() for p in model.parameters()) trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) print(fTotal params: {total_params:,}) print(fTrainable params: {trainable_params:,})方案二使用nn.LazyLinear和nn.LazyConv2d# 自动推断输入维度 model nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.LazyLinear(128), # 自动推断展平后的维度 nn.ReLU(), nn.LazyLinear(10), ) # 第一次前向传播时自动推断 x torch.randn(4, 3, 32, 32) output model(x) print(fOutput shape: {output.shape})方案三使用torchinfo库# pip install torchinfo from torchinfo import summary model nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(16 * 16 * 16, 10), ) # 打印模型摘要 summary(model, input_size(4, 3, 32, 32))完整修复代码 完整的 PyTorch 模型输入形状管理方案 涵盖形状检查、形状追踪、常见模型架构、调试工具 import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader, TensorDataset from typing import Tuple, Optional, List, Dict import math # # 形状检查和追踪工具 # class ShapeChecker: 模型形状检查和追踪工具 staticmethod def trace_shapes(model: nn.Module, input_tensor: torch.Tensor) - List[Dict]: 追踪模型每层的输入输出形状 shapes [] hooks [] def hook_fn(module, input, output): input_shape tuple(input[0].shape) if isinstance(input, tuple) else tuple(input.shape) if isinstance(output, torch.Tensor): output_shape tuple(output.shape) elif isinstance(output, tuple): output_shape tuple(tuple(o.shape) for o in output) else: output_shape str(type(output)) layer_name module.__class__.__name__ params sum(p.numel() for p in module.parameters()) shapes.append({ layer: layer_name, input_shape: input_shape, output_shape: output_shape, params: params, }) for module in model.modules(): if len(list(module.children())) 0: h module.register_forward_hook(hook_fn) hooks.append(h) with torch.no_grad(): _ model(input_tensor) for h in hooks: h.remove() return shapes staticmethod def print_summary(model: nn.Module, input_tensor: torch.Tensor): 打印模型摘要 shapes ShapeChecker.trace_shapes(model, input_tensor) print(f\n{*80})  print(f{Layer:25} {Input Shape:25} {Output Shape:25} {Params:10}) print(f{-*80}) total_params 0 for info in shapes: print(f{info[layer]:25} {str(info[input_shape]):25} f{str(info[output_shape]):25} {info[params]:10,}) total_params info[params] print(f{-*80}) print(f{Total:25} {:25} {:25} {total_params:10,}) print(f{*80}\n) staticmethod def check_input_compatibility(model: nn.Module, input_shape: Tuple[int, ...]) - bool: 检查输入形状是否与模型兼容 try: dummy_input torch.randn(*input_shape) with torch.no_grad(): _ model(dummy_input) return True except RuntimeError as e: print(f形状不兼容: {e}) return False staticmethod def get_output_shape(model: nn.Module, input_shape: Tuple[int, ...]) - Tuple[int, ...]: 获取模型的输出形状 dummy_input torch.randn(*input_shape) with torch.no_grad(): output model(dummy_input) return tuple(output.shape) # # 常见模型架构及其输入形状 # class MLPModel(nn.Module): 多层感知机 - 输入: [batch_size, input_size] def __init__(self, input_size784, hidden_size128, num_classes10): super().__init__() self.model nn.Sequential( nn.Linear(input_size, hidden_size), nn.ReLU(), nn.Dropout(0.2), nn.Linear(hidden_size, hidden_size // 2), nn.ReLU(), nn.Linear(hidden_size // 2, num_classes), ) def forward(self, x): return self.model(x) class CNNModel(nn.Module): CNN - 输入: [batch_size, channels, height, width] def __init__(self, in_channels3, num_classes10, input_size32): super().__init__() self.features nn.Sequential( nn.Conv2d(in_channels, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) # 计算展平后维度 feat_size input_size // 4 # 两次 MaxPool2d(2) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(64 * feat_size * feat_size, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes), ) def forward(self, x): x self.features(x) x self.classifier(x) return x class RNNModel(nn.Module): RNN/LSTM - 输入: [batch_size, seq_len, input_size] def __init__(self, input_size10, hidden_size20, num_layers2, num_classes3, rnn_typelstm): super().__init__() rnn_cls nn.LSTM if rnn_type lstm else nn.GRU self.rnn rnn_cls( input_sizeinput_size, hidden_sizehidden_size, num_layersnum_layers, batch_firstTrue, bidirectionalFalse, dropout0.1 if num_layers 1 else 0, ) self.fc nn.Linear(hidden_size, num_classes) def forward(self, x): # x: [batch_size, seq_len, input_size] out, _ self.rnn(x) # out: [batch_size, seq_len, hidden_size] out out[:, -1, :] # 取最后一个时间步: [batch_size, hidden_size] out self.fc(out) # [batch_size, num_classes] return out class TextClassifier(nn.Module): 文本分类器 - 输入: [batch_size, seq_len] (整数索引) def __init__(self, vocab_size10000, embed_dim128, hidden_size256, num_classes5, num_layers2): super().__init__() self.embedding nn.Embedding(vocab_size, embed_dim, padding_idx0) self.lstm nn.LSTM( input_sizeembed_dim, hidden_sizehidden_size, num_layersnum_layers, batch_firstTrue, bidirectionalTrue, dropout0.3, ) self.fc nn.Sequential( nn.Linear(hidden_size * 2, hidden_size), # *2 因为双向 nn.ReLU(), nn.Dropout(0.3), nn.Linear(hidden_size, num_classes), ) def forward(self, x): # x: [batch_size, seq_len] 整数索引 embedded self.embedding(x) # [batch_size, seq_len, embed_dim] out, _ self.lstm(embedded) # [batch_size, seq_len, hidden_size*2] out out[:, -1, :] # [batch_size, hidden_size*2] return self.fc(out) # [batch_size, num_classes] class MultiInputModel(nn.Module): 多输入模型 - 接收不同形状的输入 def __init__(self, image_channels3, text_vocab10000, num_classes10): super().__init__() # 图像分支 self.image_branch nn.Sequential( nn.Conv2d(image_channels, 32, kernel_size3, padding1), nn.ReLU(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 64), ) # 文本分支 self.text_branch nn.Sequential( nn.Embedding(text_vocab, 64), nn.LSTM(64, 64, batch_firstTrue), ) # 融合层 self.fusion nn.Sequential( nn.Linear(64 64, 128), nn.ReLU(), nn.Linear(128, num_classes), ) def forward(self, image, text): # image: [batch_size, 3, H, W] # text: [batch_size, seq_len] img_feat self.image_branch(image) # [batch_size, 64] text_feat, _ self.text_branch(text) text_feat text_feat[:, -1, :] # [batch_size, 64] combined torch.cat([img_feat, text_feat], dim1) # [batch_size, 128] return self.fusion(combined) # [batch_size, num_classes] # # 形状安全的 DataLoader # class ShapeSafeDataset(Dataset): 确保数据形状一致的 Dataset def __init__(self, data, labels, expected_shapeNone): self.data data self.labels labels self.expected_shape expected_shape def __len__(self): return len(self.data) def __getitem__(self, idx): sample self.data[idx] label self.labels[idx] # 确保形状正确 if self.expected_shape: if sample.shape ! self.expected_shape: sample sample.reshape(self.expected_shape) # 确保是 tensor if not isinstance(sample, torch.Tensor): sample torch.tensor(sample, dtypetorch.float32) if not isinstance(label, torch.Tensor): label torch.tensor(label, dtypetorch.long) return sample, label # # 使用示例 # def demo_mlp_shapes(): MLP 模型形状 print( * 60) print(示例 1: MLP 模型输入形状) print( * 60) model MLPModel(input_size784, hidden_size128, num_classes10) # 正确输入 x torch.randn(32, 784) # [batch_size, input_size] print(f输入形状: {x.shape}) output model(x) print(f输出形状: {output.shape}) # 打印模型摘要 ShapeChecker.print_summary(model, x) # 错误输入 print(尝试错误输入:) try: x_wrong torch.randn(32, 100) _ model(x_wrong) except RuntimeError as e: print(f 错误: {e}) print() def demo_cnn_shapes(): CNN 模型形状 print( * 60) print(示例 2: CNN 模型输入形状) print( * 60) model CNNModel(in_channels3, num_classes10, input_size32) # 正确输入 x torch.randn(4, 3, 32, 32) # [B, C, H, W] print(f输入形状: {x.shape}) output model(x) print(f输出形状: {output.shape}) ShapeChecker.print_summary(model, x) # 错误缺少 batch 维度 print(尝试 3D 输入缺少 batch 维度:) try: x_wrong torch.randn(3, 32, 32) _ model(x_wrong) except RuntimeError as e: print(f 错误: {e}) print() def demo_rnn_shapes(): RNN 模型形状 print( * 60) print(示例 3: RNN 模型输入形状) print( * 60) model RNNModel(input_size10, hidden_size20, num_classes3) # 正确输入: [batch_size, seq_len, input_size] x torch.randn(32, 50, 10) print(f输入形状: {x.shape}) output model(x) print(f输出形状: {output.shape}) ShapeChecker.print_summary(model, x) print() def demo_text_classifier_shapes(): 文本分类器形状 print( * 60) print(示例 4: 文本分类器输入形状) print( * 60) model TextClassifier(vocab_size10000, embed_dim128, hidden_size256, num_classes5) # 输入: [batch_size, seq_len] 整数索引 x torch.randint(0, 10000, (32, 50)) # 32个样本每个50个词 print(f输入形状: {x.shape}, dtype: {x.dtype}) output model(x) print(f输出形状: {output.shape}) ShapeChecker.print_summary(model, x) print() def demo_multi_input_shapes(): 多输入模型形状 print( * 60) print(示例 5: 多输入模型形状) print( * 60) model MultiInputModel(image_channels3, text_vocab10000, num_classes10) # 图像输入: [B, 3, H, W] image torch.randn(4, 3, 32, 32) # 文本输入: [B, seq_len] text torch.randint(0, 10000, (4, 20)) print(f图像输入形状: {image.shape}) print(f文本输入形状: {text.shape}) output model(image, text) print(f输出形状: {output.shape}) print() def demo_shape_fixing(): 形状修复示例 print( * 60) print(示例 6: 常见形状修复) print( * 60) # 1. 添加 batch 维度 single_sample torch.randn(3, 32, 32) # [C, H, W] batched single_sample.unsqueeze(0) # [1, C, H, W] print(f添加 batch 维度: {single_sample.shape} - {batched.shape}) # 2. 展平 conv_output torch.randn(4, 16, 8, 8) # [B, C, H, W] flattened conv_output.view(conv_output.size(0), -1) # [B, C*H*W] print(f展平: {conv_output.shape} - {flattened.shape}) # 3. 重塑 features torch.randn(4, 1024) reshaped features.reshape(4, 16, 8, 8) # [B, C, H, W] print(f重塑: {features.shape} - {reshaped.shape}) # 4. 交换维度 # RNN 默认: [seq_len, batch_size, features] # batch_first: [batch_size, seq_len, features] rnn_input torch.randn(50, 32, 10) # [seq, batch, feat] batch_first rnn_input.transpose(0, 1) # [batch, seq, feat] print(f交换维度: {rnn_input.shape} - {batch_first.shape}) # 5. 扩展维度 scalar torch.randn(10) # [10] expanded scalar.unsqueeze(1) # [10, 1] print(f扩展维度: {scalar.shape} - {expanded.shape}) print() def demo_dataloader_shapes(): DataLoader 形状 print( * 60) print(示例 7: DataLoader 输出形状) print( * 60) # 创建数据 X torch.randn(100, 3, 32, 32) # 100张 3x32x32 的图片 y torch.randint(0, 10, (100,)) dataset TensorDataset(X, y) # DataLoader 自动添加 batch 维度 dataloader DataLoader(dataset, batch_size16, shuffleTrue) for batch_x, batch_y in dataloader: print(fBatch X 形状: {batch_x.shape}) # [16, 3, 32, 32] print(fBatch Y 形状: {batch_y.shape}) # [16] break print() if __name__ __main__: demo_mlp_shapes() demo_cnn_shapes() demo_rnn_shapes() demo_text_classifier_shapes() demo_multi_input_shapes() demo_shape_fixing() demo_dataloader_shapes() print( * 60) print(所有示例执行完毕) print( * 60)常见陷阱与注意事项1. 始终添加 batch 维度# 单个样本需要添加 batch 维度 x torch.randn(3, 32, 32) # [C, H, W] x x.unsqueeze(0) # [1, C, H, W] # 或 x x[None, ...] # [1, C, H, W]2. 注意batch_first参数# RNN 默认: [seq_len, batch_size, input_size] lstm nn.LSTM(input_size10, hidden_size20) # batch_firstFalse # batch_firstTrue: [batch_size, seq_len, input_size] lstm nn.LSTM(input_size10, hidden_size20, batch_firstTrue)3. 卷积输出尺寸计算# Conv2d 输出尺寸: # H_out floor((H_in 2*padding - kernel_size) / stride) 1 # 展平后维度 out_channels * H_out * W_out4. 使用torchinfo检查模型from torchinfo import summary summary(model, input_size(batch_size, channels, height, width))5. Embedding 层输入是整数# Embedding 输入是整数索引不是浮点数 x torch.randint(0, vocab_size, (batch_size, seq_len)) # 正确 # x torch.randn(batch_size, seq_len) # 错误6. 动态形状使用x.size(0)# 在 forward 中使用动态 batch_size def forward(self, x): batch_size x.size(0) x x.view(batch_size, -1) # 展平保留 batch 维度 return self.fc(x)总结在 PyTorch 中正确处理模型输入形状关键要点如下理解各层的输入形状约定Linear 期望 2DConv2d 期望 4DRNN 期望 3DEmbedding 期望整数索引。始终添加 batch 维度即使单个样本也要unsqueeze(0)添加 batch 维度。使用形状追踪工具通过register_forward_hook追踪每层的输入输出形状快速定位不匹配。使用nn.LazyLinear自动推断展平后的输入维度避免手动计算错误。注意batch_first参数RNN 默认batch_firstFalse需要根据数据形状设置。正确计算卷积输出尺寸使用公式floor((input 2*padding - kernel) / stride) 1。在 forward 中使用动态形状使用x.size(0)获取 batch_size避免硬编码。使用torchinfo库类似 Keras 的model.summary()直观查看模型结构和形状。通过理解这些形状约定和调试技巧可以快速定位和解决 PyTorch 中的输入形状不匹配问题提高模型开发效率。