PyTorch 中的转置操作：深入理解 `transpose()` 和 `permute()

原创已于 2025-04-01 10:37:44 修改 · 709 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

于 2025-04-01 10:06:51 首次发布

笔记同时被 2 个专栏收录

23 篇文章

订阅专栏

ML|DL

7 篇文章

订阅专栏

PyTorch 中的转置操作：深入理解 `transpose()` 和 `permute()`

引言

在深度学习的日常工作中，我们经常需要调整张量的形状和维度顺序。无论是处理图像数据、序列数据，还是实现复杂的神经网络架构，张量维度的操作都是必不可少的技能。PyTorch 作为当前最流行的深度学习框架之一，提供了两个强大的维度转置函数：transpose() 和 permute()。这两个函数虽然都用于调整张量的维度顺序，但在功能、适用场景和性能上存在显著差异。

本文将深入探讨这两个函数的原理、使用方法、区别以及实际应用场景，帮助你全面掌握 PyTorch 中的维度转置技术。无论你是刚接触深度学习的新手，还是寻求优化模型性能的资深研究者，本文都将为你提供有价值的见解。

张量和维度的基础概念

在深入讨论转置函数前，我们需要明确理解 PyTorch 中张量和维度的基本概念。

**张量（Tensor）**是 PyTorch 中的基本数据结构，可以看作是多维数组。从数学角度看：

标量是 0 维张量（如单个数字）
向量是 1 维张量（如一维数组）
矩阵是 2 维张量（如二维数组）
更高维的数组则是更高维的张量

每个张量都有一个形状（shape），表示各个维度的大小。例如，形状为 [3, 4, 5] 的张量是一个 3 维张量，第一维大小为 3，第二维大小为 4，第三维大小为 5。

**维度（Dimension）或轴（Axis）**指的是张量的轴数。在 PyTorch 中，张量的维度从 0 开始索引。例如，对于一个 3 维张量，其维度索引为 0、1、2。

import torch

# 创建一个 3x4x5 的张量
tensor = torch.randn(3, 4, 5)
print(f"张量形状: {tensor.shape}")  # torch.Size([3, 4, 5])
print(f"张量维度数: {tensor.dim()}")  # 3

理解张量的维度是使用 transpose() 和 permute() 的基础。

`transpose()` 函数详解

功能与原理

transpose() 是 PyTorch 中用于交换张量两个维度的基本函数。它的主要作用是创建一个新的视图（view），指向原始张量的存储，但具有不同的步长（stride）。这意味着 transpose() 操作本身不会复制数据，而只是改变了数据的访问方式。

语法与参数

# 方法形式
tensor.transpose(dim0, dim1)

# 函数形式
torch.transpose(input, dim0, dim1)

参数说明：

input：需要进行维度变换的张量（函数形式需要）
dim0：要交换的第一个维度的索引
dim1：要交换的第二个维度的索引

返回值：返回一个新的张量视图，具有交换后的维度

使用示例

示例 1：二维张量转置

对于一个二维张量（矩阵），transpose() 的作用相当于交换行和列：

import torch

# 创建一个 2x3 的张量
tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(f"原始张量:\n{tensor}")
print(f"原始张量形状: {tensor.shape}")

# 使用 transpose 交换维度 0 和 1
transposed = tensor.transpose(0, 1)
print(f"转置后张量:\n{transposed}")
print(f"转置后张量形状: {transposed.shape}")

输出结果：

原始张量:
tensor([[1, 2, 3],
        [4, 5, 6]])
原始张量形状: torch.Size([2, 3])
转置后张量:
tensor([[1, 4],
        [2, 5],
        [3, 6]])
转置后张量形状: torch.Size([3, 2])

示例 2：三维张量转置

对于一个三维张量，transpose() 可以交换任意两个维度：

# 创建一个 2x3x4 的张量
tensor = torch.randn(2, 3, 4)
print(f"原始张量形状: {tensor.shape}")

# 使用 transpose 交换维度 0 和 2
transposed = tensor.transpose(0, 2)
print(f"转置后张量形状: {transposed.shape}")

输出结果：

原始张量形状: torch.Size([2, 3, 4])
转置后张量形状: torch.Size([4, 3, 2])

数据共享验证

由于 transpose() 创建的是原始数据的视图，修改转置后的张量会影响原始张量：

# 创建一个 2x3 的张量
tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
transposed = tensor.transpose(0, 1)

# 修改转置后张量的值
transposed[0, 0] = 99
print(f"修改后的原始张量:\n{tensor}")

输出结果：

修改后的原始张量:
tensor([[99,  2,  3],
        [ 4,  5,  6]])

`permute()` 函数详解

功能与原理

与 transpose() 只能交换两个维度相比，permute() 函数可以一次性重新排列张量的所有维度。它同样不会复制数据，而是创建一个新的视图。

语法与参数

# 方法形式
tensor.permute(*dims)

# 函数形式
torch.permute(input, dims)

参数说明：

input：需要进行维度变换的张量（函数形式需要）
dims：一个包含所有维度索引的元组或列表，表示新的维度顺序

返回值：返回一个具有指定顺序的维度的新张量视图

使用示例

示例 1：三维张量维度重排

import torch

# 创建一个 2x3x4 的张量
tensor = torch.randn(2, 3, 4)
print(f"原始张量形状: {tensor.shape}")

# 使用 permute 重新排列维度
permuted = tensor.permute(2, 0, 1)
print(f"重排后张量形状: {permuted.shape}")

输出结果：

原始张量形状: torch.Size([2, 3, 4])
重排后张量形状: torch.Size([4, 2, 3])

示例 2：四维张量维度重排

# 创建一个 2x3x4x5 的张量
tensor = torch.randn(2, 3, 4, 5)
print(f"原始张量形状: {tensor.shape}")

# 使用 permute 将维度顺序重排为 (3, 1, 0, 2)
permuted = tensor.permute(3, 1, 0, 2)
print(f"重排后张量形状: {permuted.shape}")

输出结果：

原始张量形状: torch.Size([2, 3, 4, 5])
重排后张量形状: torch.Size([5, 3, 2, 4])

两者的区别与联系

虽然 transpose() 和 permute() 都可以用于调整张量的维度顺序，但它们之间存在一些重要的区别：

1. 功能范围

transpose() 只能交换两个维度
permute() 可以任意重新排列所有维度

2. 实现关系

从实现上看，transpose(dim0, dim1) 可以看作是 permute() 的特例。对于 n 维张量，tensor.transpose(dim0, dim1) 等价于创建一个维度序列，其中 dim0 和 dim1 的位置互换，然后调用 permute()。

下面的代码演示了 transpose() 和等价的 permute() 操作：

import torch

tensor = torch.randn(2, 3, 4, 5)

# 使用 transpose 交换维度 1 和 3
trans = tensor.transpose(1, 3)
print(f"使用 transpose 后的形状: {trans.shape}")

# 等价的 permute 操作
dims = list(range(tensor.dim()))
dims[1], dims[3] = dims[3], dims[1]
perm = tensor.permute(dims)
print(f"使用等价 permute 后的形状: {perm.shape}")

# 验证两者结果相同
print(f"两者是否完全相同: {torch.all(trans == perm)}")

输出结果：

使用 transpose 后的形状: torch.Size([2, 5, 4, 3])
使用等价 permute 后的形状: torch.Size([2, 5, 4, 3])
两者是否完全相同: tensor(True)

3. 使用场景

当只需要交换两个维度时，transpose() 更为简洁
当需要复杂的维度重排时，permute() 是唯一选择
对于二维矩阵的简单转置，transpose() 是最常用的选择
对于高维张量的复杂变换，permute() 提供了更大的灵活性

内存布局与连续性

理解 transpose() 和 permute() 的一个关键点是张量的内存布局和连续性（contiguity）。

内存布局

在 PyTorch 中，张量在内存中是按行优先（row-major）顺序存储的，这意味着最右边的维度（最内层的循环）在内存中是相邻的。例如，对于一个形状为 [2, 3, 4] 的张量，内存中的元素顺序是：

[0,0,0], [0,0,1], [0,0,2], [0,0,3], [0,1,0], [0,1,1], ...

连续性

一个张量被称为"连续的"（contiguous），如果它的元素在内存中的存储顺序与按行优先遍历张量时访问元素的顺序一致。

当我们使用 transpose() 或 permute() 重新排列维度时，我们改变了张量的步长（stride），这可能导致张量变为非连续的。非连续张量在某些操作中可能会导致性能下降，因为 PyTorch 可能需要先将其转换为连续张量。

import torch

tensor = torch.randn(2, 3, 4)
print(f"原始张量是否连续: {tensor.is_contiguous()}")  # True

transposed = tensor.transpose(0, 2)
print(f"转置后张量是否连续: {transposed.is_contiguous()}")  # False

permuted = tensor.permute(1, 0, 2)
print(f"重排后张量是否连续: {permuted.is_contiguous()}")  # False

contiguous() 方法

我们可以使用 contiguous() 方法将非连续张量转换为连续张量：

# 将非连续张量转换为连续张量
continuous = transposed.contiguous()
print(f"转换后张量是否连续: {continuous.is_contiguous()}")  # True

需要注意的是，contiguous() 方法会创建一个新的张量，并复制数据，这可能会导致额外的内存开销。因此，只有在必要时才应该使用它。

性能考量

在深度学习中，性能往往是一个关键因素。以下是关于 transpose() 和 permute() 性能的一些考量：

1. 内存访问模式

转置操作会改变张量的步长，可能导致后续操作的内存访问不连续。不连续的内存访问可能导致缓存命中率降低，从而影响性能。

2. 连续性转换开销

如果后续操作需要连续张量，则可能需要调用 contiguous()。contiguous() 会复制数据，这可能导致显著的性能开销，特别是对于大型张量。

import torch
import time

def benchmark(func, name, iterations=1000):
    start = time.time()
    for _ in range(iterations):
        result = func()
    end = time.time()
    print(f"{name}: {(end - start) * 1000 / iterations:.4f} ms per iteration")
    return result

tensor = torch.randn(100, 200, 300, device='cuda' if torch.cuda.is_available() else 'cpu')

# 测试 transpose
benchmark(lambda: tensor.transpose(0, 2), "transpose")

# 测试 permute
benchmark(lambda: tensor.permute(2, 1, 0), "permute")

# 测试 transpose + contiguous
benchmark(lambda: tensor.transpose(0, 2).contiguous(), "transpose + contiguous")

# 测试 permute + contiguous
benchmark(lambda: tensor.permute(2, 1, 0).contiguous(), "permute + contiguous")

通常，transpose() 和 permute() 本身的性能差异不大，但如果需要连续张量，则额外的 contiguous() 调用会带来显著开销。

3. 选择建议

如果只需要交换两个维度，优先使用 transpose()
如果需要复杂的维度重排，使用 permute()
如果后续操作需要连续张量，可以考虑直接使用 permute() 后跟 contiguous()
尽量避免频繁的维度变换和连续性转换

实际应用场景

了解了 transpose() 和 permute() 的基本原理和性能考量后，让我们看看它们在实际深度学习任务中的应用场景。

1. 图像处理中的通道转换

在计算机视觉任务中，图像数据的格式转换是一个常见需求。不同的库和框架可能使用不同的格式：

PyTorch 默认使用 [batch_size, channels, height, width] (NCHW) 格式
TensorFlow 默认使用 [batch_size, height, width, channels] (NHWC) 格式
OpenCV 使用 [height, width, channels] (HWC) 格式

在这些格式之间转换时，transpose() 和 permute() 非常有用：

import torch
import numpy as np
import cv2

# 从 OpenCV 加载图像 (HWC 格式)
image = cv2.imread('image.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # OpenCV 使用 BGR，转换为 RGB

# 转换为 PyTorch 张量
tensor = torch.from_numpy(image).float()
print(f"OpenCV 图像张量形状: {tensor.shape}")  # [H, W, C]

# 转换为 PyTorch 期望的格式 (添加批次维度并调整通道位置)
tensor = tensor.permute(2, 0, 1).unsqueeze(0)
print(f"PyTorch 图像张量形状: {tensor.shape}")  # [N, C, H, W]

# 如果需要转换回 OpenCV 格式
opencv_tensor = tensor.squeeze(0).permute(1, 2, 0)
print(f"转换回 OpenCV 格式: {opencv_tensor.shape}")  # [H, W, C]

2. 序列模型中的序列处理

在自然语言处理和时间序列分析中，经常需要调整序列数据的维度：

import torch

# 假设我们有一个 [batch_size, sequence_length, features] 的序列数据
sequence = torch.randn(32, 100, 512)

# 对于某些 RNN 实现，可能需要 [sequence_length, batch_size, features] 格式
rnn_input = sequence.transpose(0, 1)
print(f"RNN 输入形状: {rnn_input.shape}")  # [100, 32, 512]

# 对于某些操作，可能需要将特征维度放在前面
feature_first = sequence.permute(2, 0, 1)
print(f"特征优先形状: {feature_first.shape}")  # [512, 32, 100]

3. 注意力机制中的矩阵运算

在 Transformer 等架构中，注意力机制涉及复杂的矩阵运算，常常需要调整维度：

import torch

# 假设我们有查询、键和值矩阵
batch_size, seq_len, d_model = 16, 50, 512
query = torch.randn(batch_size, seq_len, d_model)
key = torch.randn(batch_size, seq_len, d_model)
value = torch.randn(batch_size, seq_len, d_model)

# 计算注意力分数
# 将 key 转置以进行矩阵乘法
scores = torch.matmul(query, key.transpose(1, 2))
print(f"注意力分数形状: {scores.shape}")  # [batch_size, seq_len, seq_len]

# 多头注意力中的维度调整
num_heads = 8
head_dim = d_model // num_heads

# 将 query 重塑为多头形式
query_multi_head = query.view(batch_size, seq_len, num_heads, head_dim)
# 调整维度顺序为 [batch_size, num_heads, seq_len, head_dim]
query_multi_head = query_multi_head.permute(0, 2, 1, 3)
print(f"多头查询形状: {query_multi_head.shape}")  # [16, 8, 50, 64]