Transformer中对x.flatten(2).permute(2, 0, 1)的学习

在学习Transformer中遇到了x.flatten(2).permute(2, 0, 1)，乍一看就很头大，完全不知道什么意思。下面一起来学习一下吧。

1. permute函数

作用是改变tensor的内部维度的顺序。

最初始为permute(0, 1，2)，0、1、2依次表示通道、行、列三个概念。

permute(2, 0, 1)就表示将原来的列提到最前面，另外两个依次退后一个。

2. flatten函数

将维度展开为一维的函数。

python里的flatten(dim)表示，从第dim个维度开始展开，将后面的维度转化为一维。

也就是说，只保留dim之前的维度，其他维度的数据全都挤在dim这一维。

此处的flatten(2)表示将第2个维度及以后展开到第2维上。

3. 总体理解

先flatten再permute，即先将原来的第2维（实际是序号为三）后的维度展平到第2维，再将原来的2维（即列）提到最前面，另外两个依次退后一个。