PyTorch技术问询：如何打印网络各层输出Blob尺寸及简单层索引

嘿，我来帮你搞定这俩需求！既然已经能打印网络结构和权重尺寸了，这两个问题其实都是在遍历网络层的基础上做些针对性扩展就行，咱们分情况聊：

1. 打印网络中每个层的输出Blob尺寸 #

这里分两种常用框架给你具体实现方案：

PyTorch 实现 #

最便捷的方式是给每个运算层注册前向钩子（forward hook），在模型前向传播时自动捕获输出尺寸，还能跳过容器类模块（比如Fire、Sequential），只关注实际产生输出的层：

import torch
from torch import nn

def log_output_shape(module, input, output):
    # 打印层名称、类型和输出尺寸
    print(f"Layer: {module.__class__.__name__}, Output shape: {output.shape}")

# 假设你的模型实例是model
for name, layer in model.named_modules():
    # 只给没有子模块的"简单层"注册钩子
    if not list(layer.children()):
        layer.register_forward_hook(log_output_shape)

# 输入一个测试张量触发前向传播（尺寸根据你的模型输入调整）
test_input = torch.randn(1, 3, 224, 224)
model(test_input)

Caffe 实现 #

通过Caffe的Python接口，直接在 forward 后遍历所有Blob，关联到对应层的输出：

import caffe

# 加载网络
net = caffe.Net('deploy.prototxt', 'your_model.caffemodel', caffe.TEST)
# 准备测试输入（需匹配网络输入尺寸）
test_input = ... 
net.forward(data=test_input)

# 遍历所有层，输出对应Blob的尺寸
for layer_name in net._layer_names:
    # 获取当前层的第一个输出Blob
    output_blob = net.blobs[net.layers[layer_name].top[0]]
    print(f"Layer {layer_name} output shape: {output_blob.data.shape}")

2. 打印每个“简单层”的位置索引 #

核心是区分容器模块（如Fire、Sequential）和简单运算层（如Conv2d、ReLU、Pooling），给后者分配全局递增的索引：

PyTorch 实现 #

用递归遍历的方式，跳过容器模块，只给内部的简单层计数：

from torch import nn

# 初始化全局索引计数器
simple_layer_idx = 0

def traverse_and_index(module, parent_path=""):
    global simple_layer_idx
    for name, child in module.named_children():
        current_layer_path = f"{parent_path}.{name}" if parent_path else name
        # 判断是否为简单层：没有子模块的就是
        if not list(child.children()):
            simple_layer_idx += 1
            print(f"Simple layer index: {simple_layer_idx}, Path: {current_layer_path}, Type: {child.__class__.__name__}")
        else:
            # 递归遍历容器模块内部的子层
            traverse_and_index(child, current_layer_path)

# 对你的模型执行遍历索引
traverse_and_index(model)

Caffe 实现 #

Caffe的复合层（如Fire）内部会嵌套子层，直接遍历所有层级的层即可：

import caffe

net = caffe.Net('deploy.prototxt', 'your_model.caffemodel', caffe.TEST)
simple_layer_idx = 0

for layer in net.layers:
    # 如果是复合层（内部有子层），遍历子层并计数
    if hasattr(layer, 'layers'):
        for sub_layer in layer.layers:
            simple_layer_idx += 1
            print(f"Simple layer index: {simple_layer_idx}, Type: {sub_layer.type}")
    else:
        # 本身就是简单层，直接计数
        simple_layer_idx += 1
        print(f"Simple layer index: {simple_layer_idx}, Type: {layer.type}")

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