python - PyTorch flatten 不保持批量大小

Question

在 Keras 中，使用 Flatten()层保留批量大小。例如，如果 Flatten 的输入形状是 (32, 100, 100) , 在 Keras Flatten 的输出是 (32, 10000) ，但在 PyTorch 中是 320000 .为什么会这样？

Answer 1

正如 OP 在他们的回答中已经指出的那样，张量操作不会默认考虑批量维度。您可以使用 torch.flatten() 或 Tensor.flatten() 与 start_dim=1在批量维度之后开始展平操作。

或者，从 PyTorch 1.2.0 开始，您可以定义 nn.Flatten() 模型中的图层默认为 start_dim=1 .