【PyTorch】 torch.flatten()与nn.Flatten()的区别

问题

torch.flatten()与nn.Flatten()都可以实现展开Tensor，那么二者的区别是什么呢？

方法

经过查阅相关资料，发现二者主要区别有：
(1) 默认的dim不同，torch.flatten()默认的dim=0，而nn.Flatten()默认的dim=1，例如输入数据的尺寸是[3,1,4,4]，经过torch.flatten()展开后的尺寸变为[48]，而经过nn.Flatten()后得到的结果是[3, 16];
(2) nn.Flatten是一个类，而torch.flatten()则是一个函数。

相关实验代码如下所示：

import torch
from torch import nn

'''torch.flatten与nn.Flatten的区别
- torch.flatten()的dim默认是从0维开始;
- nn.Flatten()的dim默认是从1维开始;
'''

if __name__ ==