pytorch——flatten()
Posted ^_^|
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了pytorch——flatten()相关的知识，希望对你有一定的参考价值。
flatten()函数的作用是将tensor铺平成一维
torch.flatten(input, start_dim=0, end_dim=- 1) → Tensor
- input (Tensor) – the input tensor.
- start_dim (int) – the first dim to flatten
- end_dim (int) – the last dim to flatten
start_dim和end_dim构成了整个你要选择铺平的维度范围
下面举例说明
x = torch.tensor([[1,2], [3,4], [5,6]])
x = x.flatten(0)
x
------------------------
tensor([1, 2, 3, 4, 5, 6])
对于图片数据，我们往往期望进入fc层的维度为(channels, N)这样
x = torch.tensor([[[1,2],[3,4]], [[5,6],[7,8]]])
x = x.flatten(1)
x
-------------------------
tensor([[1, 2],
[3, 4],
[5, 6]])
以上是关于pytorch——flatten()的主要内容，如果未能解决你的问题，请参考以下文章
pytorch——torch.flatten() 和 torch.nn.Flatten()
PyTorch torch.flatten()与nn.Flatten()的区别
PyTorch torch.flatten()与nn.Flatten()的区别