PyTorch运行自定义EmbeddingNet代码触发输入通道不匹配RuntimeError
问题根因
报错的核心是两处参数配置与实际输入不匹配:
- 网络首层卷积
nn.Conv2d(1, 32, 5)的输入通道数设为1,是默认适配单通道灰度图的配置,但实际喂入网络的是3通道RGB图像(输入张量形状为[256, 3, 256, 256],对应批量大小256、3个颜色通道、高宽256像素),通道维度对齐失败直接触发报错 - 原网络整体是为MNIST数据集的2828灰度图设计的,除通道数错配外,输入图像尺寸从2828提升到256*256后,卷积网络输出的特征图尺寸也会和原全连接层的输入维度预设不匹配,仅修改通道数会触发下一个维度不匹配的报错
修正方案
- 调整首层卷积的输入通道参数,将
nn.Conv2d的第一个参数从1改为3,匹配RGB图像的通道数 - 解决不同输入尺寸带来的特征图维度不匹配问题,推荐两种适配方式:
- 方式一:在卷积网络末尾添加自适应平均池化层
nn.AdaptiveAvgPool2d((4, 4)),无论卷积层输出的特征图尺寸多大,都会统一池化为4*4大小,完全匹配原全连接层nn.Linear(64 * 4 * 4, 256)的输入维度,无需额外计算尺寸,适配性更强 - 方式二:根据256*256的实际输入尺寸,手动计算两层卷积+池化后的特征图尺寸,修改全连接层第一个Linear层的输入参数
- 方式一:在卷积网络末尾添加自适应平均池化层
修正后可直接运行的代码
import torch.nn as nn import torch.nn.functional as F class EmbeddingNet(nn.Module): def __init__(self): super(EmbeddingNet, self).__init__() self.convnet = nn.Sequential( # 调整首层卷积输入通道,匹配RGB图像 nn.Conv2d(3, 32, 5), nn.PReLU(), nn.MaxPool2d(2, stride=2), nn.Conv2d(32, 64, 5), nn.PReLU(), nn.MaxPool2d(2, stride=2), # 添加自适应池化统一输出特征图尺寸 nn.AdaptiveAvgPool2d((4, 4)) ) self.fc = nn.Sequential(nn.Linear(64 * 4 * 4, 256), nn.PReLU(), nn.Linear(256, 256), nn.PReLU(), nn.Linear(256, 2) ) def forward(self, x): output = self.convnet(x) output = output.view(output.size()[0], -1) output = self.fc(output) return output def get_embedding(self, x): return self.forward(x) class EmbeddingNetL2(EmbeddingNet): def __init__(self): super(EmbeddingNetL2, self).__init__() def forward(self, x): output = super(EmbeddingNetL2, self).forward(x) output /= output.pow(2).sum(1, keepdim=True).sqrt() return output def get_embedding(self, x): return self.forward(x)
验证提示:修改完成后,可以先构造和实际输入形状一致的张量做前向测试,比如执行
model = EmbeddingNet(); test_tensor = torch.randn(256,3,256,256); print(model(test_tensor).shape),确认无报错后再接入正式训练流程。
内容的提问来源于stack exchange,提问作者Beta
相关产品推荐
相关产品推荐

