求PyTorch/TensorFlow/Numpy中CNN输出形状计算的内置函数

我每次根据超参数计算卷积神经网络（CNN）的输出形状时，都得重复写相同的独立函数，有时候还要写单元测试，真的烦透了。

PyTorch中nn.Conv3d的输出形状计算公式：

请问有没有能实现上述公式的现成库函数？优先考虑PyTorch、TensorFlow或Numpy生态的库。

以下是我自己为PyTorch项目（Python 3.10+）写的实现函数，但我更想直接导入现成的：

def conv_conversion(
    in_shape: tuple[int, ...],
    kernel_size: int | tuple[int, ...],
    padding: int | tuple[int, ...] = 0,
    dilation: int | tuple[int, ...] = 1,
    stride: int | tuple[int, ...] = 1,
) -> tuple[int, ...]:
    """计算与nn.Conv默认行为一致的卷积层输出形状"""

    def to_tuple(value: int | tuple[int, ...]) -> tuple[int, ...]:
        return (value,) * len(in_shape) if isinstance(value, int) else value

    k, p = to_tuple(kernel_size), to_tuple(padding)
    dil, s = to_tuple(dilation), to_tuple(stride)
    return tuple(
        int((in_shape[i] + 2 * p[i] - dil[i] * (k[i] - 1) - 1) / s[i] + 1)
        for i in range(len(in_shape))
    )

现成解决方案 #

PyTorch 内置方式 #

不用自己写公式，PyTorch的Conv1d/Conv2d/Conv3d层可以直接通过构造输入张量推导输出形状，完全匹配内部计算逻辑：

import torch
import torch.nn as nn

def get_conv_out_shape(in_shape, conv_layer):
    # 构造一个符合输入形状的虚拟张量（batch_size设为1）
    dummy_input = torch.randn(1, *in_shape)
    with torch.no_grad():
        output = conv_layer(dummy_input)
    return output.shape[1:]  # 去掉batch维度，返回输出的通道+空间形状

# 示例：计算Conv3D的输出形状
conv3d = nn.Conv3d(
    in_channels=3,
    out_channels=16,
    kernel_size=3,
    padding=1,
    stride=2,
    dilation=1
)
input_shape = (3, 32, 32, 32)  # PyTorch格式：(通道数, 深度, 高度, 宽度)
print(get_conv_out_shape(input_shape, conv3d))  # 输出 torch.Size([16, 16, 16, 16])

第三方工具库 #

torchinfo（PyTorch生态）

除了计算形状，还能统计模型参数、内存占用，适合复杂模型：

from torchinfo import summary

conv3d = nn.Conv3d(3, 16, 3, padding=1, stride=2)
# input_size格式：(batch_size, 通道数, 深度, 高度, 宽度)
summary(conv3d, input_size=(1, 3, 32, 32, 32))

输出结果会清晰展示每一层的输出形状。

TensorFlow/Keras

Keras的Conv3D层自带compute_output_shape方法，直接计算：

from tensorflow.keras.layers import Conv3D

conv3d = Conv3D(16, 3, padding='same', strides=2, input_shape=(32, 32, 32, 3))
# input_shape格式：(None, 深度, 高度, 宽度, 通道数)
output_shape = conv3d.compute_output_shape(input_shape=(None, 32, 32, 32, 3))
print(output_shape)  # 输出 TensorShape([None, 16, 16, 16, 16])

纯Numpy生态库 #

如果要脱离深度学习框架，mlxtend的工具函数可以直接计算：

from mlxtend.utils import conv_output_shape

input_shape = (32, 32, 32)
output_shape = conv_output_shape(
    input_shape=input_shape,
    kernel_size=3,
    padding=1,
    stride=2,
    dilation=1
)
print(output_shape)  # 输出 (16, 16, 16)

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