TensorFlow新手求教：如何查看网络层的输出尺寸？

嘿，作为TensorFlow新手，其实你不用太担心——现在TensorFlow和Keras已经深度整合了，很多你熟悉的Keras用法在TensorFlow里完全适用！下面给你几种实用的方法来查看网络层的输出尺寸：

1. 直接使用model.summary()（最简便，和Keras用法一致） #

如果你是用TensorFlow的高阶API（tf.keras）构建的模型（比如Sequential或者Functional API），直接调用model.summary()就能看到每一层的输出维度，和你之前用Keras时的体验完全一样。

示例代码：

import tensorflow as tf

# 构建一个简单的CNN示例模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 打印模型结构及各层输出尺寸
model.summary()

输出里的Output Shape列就是你需要的各层输出尺寸，其中第一个维度是批量大小（None表示可以适配任意批量）。

2. 传入样本张量查看单一层的输出形状 #

如果你想单独验证某一层的输出尺寸，或者在构建模型的过程中一步步检查，可以创建一个符合输入形状的样本张量，传入该层后查看输出的shape属性。

示例代码：

import tensorflow as tf

# 定义一个单独的卷积层
conv_layer = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1))

# 创建一个模拟输入的张量（批量大小设为1，和输入形状匹配）
sample_input = tf.random.normal((1, 28, 28, 1))

# 得到该层的输出
layer_output = conv_layer(sample_input)

# 打印输出形状
print(f"该卷积层的输出形状: {layer_output.shape}")

3. 自定义模型中打印输出形状（低级API场景） #

如果你是用tf.Module自定义模型（TensorFlow低级API场景），可以在模型的__call__方法中直接打印各层的输出形状，或者同样通过传入样本输入查看。

示例代码：

import tensorflow as tf

class CustomCNN(tf.Module):
    def __init__(self):
        super().__init__()
        self.conv_layer = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')
        self.pool_layer = tf.keras.layers.MaxPooling2D((2, 2))
    
    @tf.function
    def __call__(self, inputs):
        x = self.conv_layer(inputs)
        print(f"卷积层输出形状: {x.shape}")
        x = self.pool_layer(x)
        print(f"池化层输出形状: {x.shape}")
        return x

# 实例化模型并传入样本输入
my_model = CustomCNN()
sample_input = tf.random.normal((1, 28, 28, 1))
my_model(sample_input)

总的来说，最推荐的还是model.summary()，它能快速帮你理清整个模型的结构和各层输出维度，完全符合你之前使用Keras的习惯~

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