理解pytorch系列:transpose是怎么实现的

在PyTorch中,transpose()是一种操作,它交换张量中两个指定维度的位置。实现这一点的关键在于不实际移动数据,而是通过改变张量的元数据(包括步长(stride)和尺寸(size))来达到效果。

举例来说,假设我们有一个形状为(3, 4)的二维张量,其内存布局为行优先(row-major)即C风格的。当我们对这个张量执行transpose(0, 1)操作时,我们期望该张量行变成列,列变成行,即得到一个形状为(4,3)的新视图。

这是通过以下步骤完成的:

  1. 改变尺寸 :改变size元数据,使得原本第一个维度(行)的大小与第二个维度(列)的大小交换。

  2. 改变步长 :步长(stride)是一个数组,指示了在每个维度上移动一个元素需要跳过的内存位置数。执行transpose()时,交换了两个维度的步长。在行优先存储的张量中,行的步长通常比列的步长大。

  3. 不移动数据:实际上数据并没有在内存中移动,只是改变了在这块内存空间上的解释方式。

以下是一个简单的示例:

python 复制代码
import torch

# 创建一个 3x4 的张量
x = torch.arange(12).view(3, 4)
print("Original tensor:")
print(x)
# 输出:
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]])

# 现在使用 transpose 来交换两个维度
y = x.transpose(0, 1)
print("\nTransposed tensor:")
print(y)
# 输出:
# tensor([[ 0,  4,  8],
#         [ 1,  5,  9],
#         [ 2,  6, 10],
#         [ 3,  7, 11]])

在这个例子中,x的形状发生了变化,但它的内存布局没有改变。通过调整步长和大小,transpose()创建了一个新的张量视图。

在PyTorch的底层C++实现中,同样接口会调用ATen库(张量操作库,是PyTorch的核心)中的对应函数,ATen函数会修改张量对象所关联的元数据以实现transpose()操作。这意味着实际的CPU或GPU中的数据不会因为transpose()操作而移动。这种"懒惰"操作提高了性能,特别是对于大型张量,因为它避免了不必要的数据拷贝。

相关推荐
启观川4 小时前
数据结构与算法 -第 3 章 常用算法-动态规划
数据结构·笔记·python·算法
码爸4 小时前
排序算法介绍
python·算法·排序算法
sm_926787054 小时前
RFID 标签打印的技术实现要点与二次开发实践
java·大数据·前端·c++·编辑器
by209994 小时前
C++ 多态:同一个基类接口,调用不同的函数
开发语言·c++·经验分享·多态
一木 之林4 小时前
Dify学习笔记 01 · 平台入门与实战:部署、五类应用、RAG与MCP(第1-16集)
人工智能·笔记·学习·计算机视觉
colinu4 小时前
C++ 异步任务与线程生命周期:`std::async`、`future`、`thread`、`jthread` 与协作式退出
java·jvm·c++
weixin_448119945 小时前
Datawhale Easy Data × AI:构建知识与记忆驱动的 Agent笔记8
人工智能·笔记·react.js
动词ing5 小时前
【题目练习】算法 组合总和||
算法
wuminyu5 小时前
ForkJoinPool窃取线程在扫描其他WorkQueue时的随机采样算法与Phase线程挂起唤醒机制解析
java·linux·c语言·jvm·c++
陈年老古董5 小时前
Hive 函数学习笔记(下):窗口函数与经典案例
hive·笔记·学习