Python PyTorch Flatten用法及代码示例

本文简要介绍python语言中 torch.nn.Flatten 的用法。

用法:

class torch.nn.Flatten(start_dim=1, end_dim=- 1)

参数：

• start_dim-第一个 dim 变平(默认值 = 1)。

• end_dim-最后 dim 变平(默认 = -1)。

将连续的暗淡范围展平为张量。与 Sequential 一起使用。

形状：

• 输入：(*, S_{\text{start}},..., S_{i}, ..., S_{\text{end}}, *),' 其中 S_{i} 是维度 i 的大小，* 表示任意数量的维度，包括无维度。

• 输出：(*, \prod_{i=\text{start}}^{\text{end}} S_{i}, *)。

例子：：

>>> input = torch.randn(32, 1, 5, 5)
>>> m = nn.Sequential(
>>>     nn.Conv2d(1, 32, 5, 1, 1),
>>>     nn.Flatten()
>>> )
>>> output = m(input)
>>> output.size()
torch.Size([32, 288])