python - torch.flatten() 和 nn.Flatten() 之间的区别

Question

torch.flatten() 和 torch.nn.Flatten() 有什么区别？

Answer 1

扁平化在 PyTorch 中以三种形式提供

作为张量方法(oop 风格) torch.Tensor.flatten 直接应用于张量:x.flatten() .

作为函数(函数形式) torch.flatten 申请为:torch.flatten(x) .

作为一个模块(层 nn.Module) nn.Flatten() .通常用于模型定义。

这三个都是相同的并且共享相同的实现，唯一的区别是 nn.Flatten有 start_dim设置为 1默认情况下避免展平第一个轴(通常是批处理轴)。而其他两个从 axis=0 变平至 axis=-1 - 即整个张量 - 如果没有给出参数。