PyTorch torch.flatten()与nn.Flatten()的区别
Posted 算法与编程之美
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了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__ == '__main__':
x = torch.rand(3, 1, 4, 4)
m = nn.Flatten()
out1 = m(x)
print(f'nn.Flatten(x) size is: out1.size()')
out2 = torch.flatten(x)
print(f'torch.flatten(x) size is: out2.size()')
结语
本文主要介绍了nn.Flatten()与torch.flatten()二者的区别,深入理解二者细微的区别方能更好的应用。
以上是关于PyTorch torch.flatten()与nn.Flatten()的区别的主要内容,如果未能解决你的问题,请参考以下文章
pytorch——torch.flatten() 和 torch.nn.Flatten()
PyTorch基础(15)-- torch.flatten()方法
PyTorch基础(15)-- torch.flatten()方法
PyTorch基础(15)-- torch.flatten()方法