如何在Flux.jl中查看模型输入维度及解决ResNet推理维度不匹配问题

你遇到的维度不匹配报错本质是Flux框架下2D卷积层的输入要求限制：Flux默认的2D卷积输入张量为4维结构，顺序为图像高度×图像宽度×通道数×批次大小。你提取的单张图像imgs[:, :, :, 2]是3维张量（仅包含高、宽、通道三个维度），缺少了必填的批次维度，和卷积核的4维结构不匹配，因此触发报错。

推理阶段必须和训练阶段保持输入张量的维度数量（秩）一致，仅批次大小的数值可以调整，单张图像推理也需要保留大小为1的批次维度，不能直接省略该维度。

可以通过以下3种方式确认模型要求的输入格式：

• 打印模型结构：直接执行print(model)或者display(model)，可以看到第一层卷积层的权重参数维度，格式为(卷积核高度, 卷积核宽度, 输入通道数, 输出通道数)，ResNet的首层卷积权重通常为(7,7,3,64)，由此可反推输入需要4维张量，第三维为3通道，第四维为批次维度。
• 直接查询首层权重尺寸：执行size(model.layers[1].weight)，返回结果的第三个值为要求的输入通道数，结合2D卷积的特性，可确定输入总维度为4（空间2维+通道1维+批次1维）。
• 测试张量验证：构造随机张量传入模型测试，例如model(rand(Float32, 224,224,3,1))，如果可以正常返回结果，说明该尺寸符合输入要求，若报错可根据提示调整维度。

你当前使用的cat方法可以正常生效，Flux也提供了更简洁的API实现批次维度添加：

img1 = Flux.unsqueeze(imgs[:, :, :, 2], 4)

该代码直接在第4维插入大小为1的批次维度，可读性更高。

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