torch.flatten传参

torch.flatten函数用于将输入张量展平为一维。它接受两个参数start_dim和end_dim，用于指定要展平的维度范围[^1]。

1. start_dim表示要开始展平的维度索引，end_dim表示要结束展平的维度索引。展平的范围是从start_dim到end_dim（不包括end_dim）。
2. 如果start_dim和end_dim相同，则只展平该维度。
3. 如果end_dim大于输入张量的维度数，则会引发错误。

以下是几个示例：
1. 对于输入张量shape为(100, 1, 28, 28)，使用torch.flatten(start_dim=1, end_dim=3)将第二维到第四维展平：

import torch

input_tensor = torch.randn(100, 1, 28, 28)
output_tensor = torch.flatten(input_tensor, start_dim=1, end_dim=3)
print(output_tensor.shape)  # 输出：torch.Size([100, 784])

2. 对于输入张量shape为(100, 1, 28, 28)，使用torch.flatten(start_dim=1, end_dim=2)将第二维展平：

import torch

input_tensor = torch.randn(100, 1, 28, 28)
output_tensor = torch.flatten(input_tensor, start_dim=1, end_dim=2)
print(output_tensor.shape)  # 输出：torch.Size([100, 56])

3. 对于输入张量shape为(100, 1, 28, 28)，使用torch.flatten(start_dim=1, end_dim=1)将第二维展平：

import torch

input_tensor = torch.randn(100, 1, 28, 28)
output_tensor = torch.flatten(input_tensor, start_dim=1, end_dim=1)
print(output_tensor.shape)  # 报错：RuntimeError: start_dim=1 exceeds dimension limit of 1