如何通过PyTorch模型对象自动获取CNN所需输入通道数？

核心思路 #

CNN的输入通道数由第一个卷积层的in_channels属性直接决定，因此我们只需要递归遍历模型的所有子模块，找到第一个具备该属性的卷积层即可——这是模型结构的固有参数，完全不需要外部输入信息。

优雅实现代码 #

import torch.nn as nn

def input_depth(network: nn.Module) -> int:
    # 递归遍历模型所有子模块（包括嵌套结构）
    for module in network.modules():
        # 匹配所有带in_channels属性的卷积类层（Conv1d/Conv2d/Conv3d等）
        if hasattr(module, 'in_channels'):
            return module.in_channels
    # 极端情况：输入模型无卷积层（不符合CNN定义，抛出异常）
    raise ValueError("Input network is not a valid CNN (no convolution layers found)")

关键细节说明 #

• modules()方法会递归遍历模型的所有层级结构（比如嵌套在Sequential、自定义Module中的子层），确保不会遗漏网络最前端的卷积层。
• 直接读取in_channels是最可靠的方式：该参数是卷积层初始化时就固定的固有属性，无需通过输入张量推断形状，完全符合“仅基于模型对象”的要求。
• 兼容自定义层：如果你的自定义卷积类也定义了in_channels属性，这个函数同样可以正常工作。

测试示例 #

# 测试标准CNN结构
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3)
        
model = SimpleCNN()
print(input_depth(model))  # 输出3，对应RGB图像输入

# 测试嵌套结构的CNN
nested_model = nn.Sequential(
    nn.Sequential(nn.Conv2d(1, 8, 3)),
    nn.Conv2d(8, 16, 3)
)
print(input_depth(nested_model))  # 输出1，对应灰度图像输入

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