点点工程机械
116.25M · 2026-04-01
import torch
# 创建一个 2x3 的张量
x = torch.tensor([[1, 2, 3],
[4, 5, 6]])
print(x.shape) # torch.Size([2, 3])
print(x.stride()) # (3, 1)
print(x.storage()) # [1, 2, 3, 4, 5, 6](底层是一维数组)
惊人的发现: 虽然我们看到的是 2D 矩阵,但底层存储的是 1D 数组!
这引出了一个关键问题:Tensor 如何将多维的逻辑视图映射到一维的物理存储?
Storage 是 PyTorch 中真正存储数据的地方,它是一个一维的连续内存块。
x = torch.tensor([1, 2, 3, 4, 5, 6])
print(x.storage())
# 输出:
# 1
# 2
# 3
# 4
# 5
# 6
# [torch.storage.TypedStorage(dtype=torch.int64, device=cpu) of size 6]
关键特性:
Shape(形状) 定义了 Tensor 的维度和每个维度的大小。
x = torch.arange(12) # shape: (12,)
# 同一个 Storage,不同的 Shape
y = x.view(3, 4) # shape: (3, 4) - 3行4列
z = x.view(2, 2, 3) # shape: (2, 2, 3) - 2个2x3矩阵
重要: Shape 只是"元数据",改变 Shape 不需要移动数据。
Stride(步长) 是理解 Tensor 内存布局的核心概念。
定义: Stride 告诉你,在某个维度上前进一步,需要在 Storage 中跳过多少个元素。
x = torch.tensor([[1, 2, 3],
[4, 5, 6]])
print(x.shape) # (2, 3) - 2行3列
print(x.stride()) # (3, 1)
Stride (3, 1) 的含义:
内存布局可视化:
Storage: [1, 2, 3, 4, 5, 6]
↑ ↑
[0,0] [1,0]
stride=0×3+0×1=0
stride=1×3+0×1=3
访问 x[i, j] 的内存地址 = base_offset + i × stride[0] + j × stride[1]
= 0 + i × 3 + j × 1
x[0, 0] → storage[0] = 1
x[0, 1] → storage[1] = 2
x[1, 0] → storage[3] = 4
x[1, 2] → storage[5] = 6
形式化定义:
Storage S: 一维数组 [s₀, s₁, s₂, ..., sₙ₋₁]
Tensor T: 通过投影函数 f 定义的多维视图
- Shape: (d₀, d₁, d₂, ..., dₖ)
- Stride: (σ₀, σ₁, σ₂, ..., σₖ)
- Offset: o
投影函数 f: 多维索引 → 一维索引
f(i₀, i₁, i₂, ..., iₖ) = o + i₀×σ₀ + i₁×σ₁ + i₂×σ₂ + ... + iₖ×σₖ
因此:
T[i₀, i₁, ..., iₖ] = S[f(i₀, i₁, ..., iₖ)]
三维示例:
x = torch.arange(24).view(2, 3, 4)
print(x.stride()) # (12, 4, 1)
# 访问 x[1, 2, 3]:
# storage_index = 1×12 + 2×4 + 3×1 = 12 + 8 + 3 = 23
print(x[1, 2, 3]) # tensor(23)
x = torch.arange(12)
print(f"x 的 Storage 指针: {x.storage().data_ptr()}")
# 创建多个不同的投影
y = x.view(3, 4)
z = x.view(2, 6)
w = x.view(2, 3, 2)
# 验证:它们共享同一个 Storage
print(f"y 的 Storage 指针: {y.storage().data_ptr()}")
print(f"z 的 Storage 指针: {z.storage().data_ptr()}")
print(f"w 的 Storage 指针: {w.storage().data_ptr()}")
# 所有指针相同!
# 修改任何一个会影响所有
x[0] = 999
print(f"x[0] = {x[0]}") # 999
print(f"y[0,0] = {y[0, 0]}") # 999
print(f"z[0,0] = {z[0, 0]}") # 999
可视化:同一个 Storage 的多种投影
Storage [0, 1, 2, 3, 4, 5, ..., 11]
(一维数组,永远不变)
│
┌───────────────┼───────────────┐
│ │ │
↓ ↓ ↓
投影 1 投影 2 投影 3
shape=(3,4) shape=(2,6) shape=(2,3,2)
stride=(4,1) stride=(6,1) stride=(6,2,1)
小结: 现在你理解了 Tensor 的物理本质:一个一维 Storage + 一套投影规则。
但你可能会问:为什么要这样设计?为什么不直接把数据和形状存在一起?
这就引出了 Part 2...
假设我们用传统的方式设计 Tensor:
# 传统设计:数据和形状紧密绑定
class NaiveTensor:
def __init__(self, data, shape):
self.data = data # 2D 数组 [[1,2,3], [4,5,6]]
self.shape = shape # (2, 3)
def transpose(self):
# 必须创建新数组!
new_data = [[self.data[j][i] for j in range(self.shape[0])]
for i in range(self.shape[1])]
return NaiveTensor(new_data, (self.shape[1], self.shape[0]))
问题: 每次转置都要复制数据!
性能测试:
import time
import torch
# 测试:1000×1000 矩阵转置
x = torch.randn(1000, 1000)
# 方法 1:PyTorch 的 transpose(只改元数据)
start = time.perf_counter()
y = x.transpose(0, 1)
time_pytorch = (time.perf_counter() - start) * 1000
# 方法 2:手动复制(模拟传统设计)
start = time.perf_counter()
y_copy = x.t().contiguous().clone()
time_copy = (time.perf_counter() - start) * 1000
print(f"PyTorch transpose: {time_pytorch:.4f} ms") # < 0.001 ms
print(f"手动复制: {time_copy:.4f} ms") # ~ 10 ms
print(f"加速比: {time_copy / time_pytorch:.0f}x")
# 快 10000 倍!
核心思想: 将"物理存储"和"逻辑视图"分离。
传统设计:
Tensor = 数据 + 形状
→ 改变形状 = 复制数据
PyTorch 设计:
Tensor = Storage(物理) + (Shape, Stride)(逻辑)
→ 改变形状 = 只改元数据
优势 1:零拷贝操作
x = torch.randn(1000, 1000)
# x: shape=(1000, 1000), stride=(1000, 1)
# 含义:1000×1000 的矩阵,Storage 有 100万个元素
# 所有这些操作都是瞬间完成,不复制数据:
y = x.transpose(0, 1)
# y: shape=(1000, 1000), stride=(1, 1000) ← 只交换了 stride
# 含义:转置矩阵,行列互换
# 访问 y[i,j] = storage[i×1 + j×1000] = x[j,i]
# Storage 完全没动!
z = x[::2, ::2]
# z: shape=(500, 500), stride=(2000, 2) ← stride 变成 2 倍
# 含义:每隔一行、每隔一列采样(下采样)
# 访问 z[i,j] = storage[i×2000 + j×2] = x[2i, 2j]
# Storage 完全没动!
w = x.unsqueeze(0)
# w: shape=(1, 1000, 1000), stride=(1000000, 1000, 1)
# 含义:在第 0 维插入大小为 1 的维度(增加批次维度)
# 访问 w[0,i,j] = storage[0×1000000 + i×1000 + j×1] = x[i,j]
# Storage 完全没动!
v = x.expand(10, 1000, 1000)
# v: shape=(10, 1000, 1000), stride=(0, 1000, 1) ← 注意第 0 维 stride=0!
# 含义:虚拟复制 10 份,看起来有 10 个矩阵,但实际内存只有 1 份
# 访问 v[i,j,k] = storage[i×0 + j×1000 + k×1] = x[j,k]
# 无论 i 是多少,都访问同一个位置!
# Storage 完全没动!
优势 2:内存共享
x = torch.randn(1000, 1000)
# x 占用约 8MB 内存(1000×1000×8字节)
y = x.transpose(0, 1)
# y 不占用额外内存!仍然是 8MB
# 验证:它们共享同一块内存
print(x.storage().data_ptr() == y.storage().data_ptr()) # True
# data_ptr() 返回 Storage 的内存地址,地址相同说明是同一块内存
# 因为共享内存,修改 y 会影响 x
y[0, 0] = 999
print(x[0, 0]) # 999
# 为什么?因为 y[0,0] 和 x[0,0] 访问的是 storage[0],同一个位置!
# 同理,修改 x 也会影响 y
x[0, 1] = 888
print(y[1, 0]) # 888
# 为什么?因为 x[0,1]=storage[1],y[1,0]=storage[1],都是同一个位置!
优势 3:延迟物化(Lazy Materialization)
x = torch.randn(1000, 1000)
# Storage: 100万个元素
# shape=(1000, 1000), stride=(1000, 1)
# ===== 第 1 步:转置 =====
y = x.transpose(0, 1)
# 瞬间完成(< 0.001ms)
# 只改了元数据:shape=(1000, 1000), stride=(1, 1000)
# Storage 还是原来的 100万个元素,位置没变!
print(f"y 和 x 共享内存: {y.storage().data_ptr() == x.storage().data_ptr()}") # True
# ===== 第 2 步:隔列采样 =====
z = y[:, ::2]
# 瞬间完成(< 0.001ms)
# 只改了元数据:shape=(1000, 500), stride=(1, 2000)
# Storage 还是原来的 100万个元素,位置没变!
print(f"z 和 x 共享内存: {z.storage().data_ptr() == x.storage().data_ptr()}") # True
# ===== 第 3 步:增加批次维度 =====
w = z.unsqueeze(0)
# 瞬间完成(< 0.001ms)
# 只改了元数据:shape=(1, 1000, 500), stride=(500000, 1, 2000)
# Storage 还是原来的 100万个元素,位置没变!
print(f"w 和 x 共享内存: {w.storage().data_ptr() == x.storage().data_ptr()}") # True
# 此时:做了 3 个操作,但一次数据复制都没有发生!
# w 现在是一个"不连续"的 Tensor,访问它会跳来跳去
# ===== 第 4 步:真正需要连续内存时 =====
w_cont = w.contiguous()
# 这里才真正复制数据(~ 2ms)
# 创建新的 Storage,按照 w 的逻辑顺序重新排列数据
# w_cont: shape=(1, 1000, 500), stride=(500000, 500, 1) ← 标准 stride
print(f"w_cont 和 x 共享内存: {w_cont.storage().data_ptr() == x.storage().data_ptr()}") # False
# 为什么需要 contiguous?
# 因为某些操作(如 view)要求连续的内存布局
# 现在可以安全地 reshape 了:
result = w_cont.view(1, -1) # 成功!shape=(1, 500000)
类比理解:
想象你有一本 1000 页的书(Storage),你想:
这就是"延迟":你一直在操作原书,没有复制任何一页!
但如果你要: 4. 把选中的页装订成新书(contiguous)→ 现在才需要复印机!
性能对比:
import time
x = torch.randn(1000, 1000)
# 测试延迟物化的性能优势
start = time.perf_counter()
y = x.transpose(0, 1)
z = y[:, ::2]
w = z.unsqueeze(0)
time_lazy = (time.perf_counter() - start) * 1000
start = time.perf_counter()
w_cont = w.contiguous()
time_materialize = (time.perf_counter() - start) * 1000
print(f"3 个 View 操作: {time_lazy:.4f} ms") # < 0.01 ms
print(f"最终物化(复制数据): {time_materialize:.4f} ms") # ~ 2 ms
print(f"物化慢了: {time_materialize / time_lazy:.0f}x")
关键洞察:
| 操作 | 传统设计 | PyTorch 设计 | 加速比 |
|---|---|---|---|
transpose() | 复制数据(10ms) | 只改 stride(< 0.001ms) | 10000x |
[::2] 切片 | 复制数据 | 只改 stride 和 offset | 1000x |
unsqueeze() | 复制并扩展 | 只改 shape | 无限 |
expand() | 复制多份 | stride 设为 0 | 无限 |
小结: 分离式设计让 PyTorch 可以实现零拷贝的视图操作,性能提升数千倍。
但这引出了新问题:如何通过只改变 Stride 来实现 transpose?Stride 到底是如何工作的?
这就是 Part 3 要解答的...
View 是对同一个 Storage 的不同"解释方式"。
核心原则:
x = torch.tensor([[1, 2, 3],
[4, 5, 6]])
print("原始:")
print("Shape:", x.shape) # (2, 3)
print("Stride:", x.stride()) # (3, 1)
print("Storage:", list(x.storage())) # [1, 2, 3, 4, 5, 6]
# 转置
y = x.transpose(0, 1)
print("n转置后:")
print("Shape:", y.shape) # (3, 2)
print("Stride:", y.stride()) # (1, 3) ← 注意 stride 变了!
print("Storage:", list(y.storage())) # [1, 2, 3, 4, 5, 6] ← Storage 没变!
关键发现: transpose 并没有移动数据,只是交换了 Stride!
可视化理解:
原始 x (shape=(2,3), stride=(3,1)):
Storage: [1, 2, 3, 4, 5, 6]
投影规则: x[i,j] = storage[i×3 + j×1]
x[0,0] = storage[0×3 + 0×1] = storage[0] = 1
x[0,1] = storage[0×3 + 1×1] = storage[1] = 2
x[1,0] = storage[1×3 + 0×1] = storage[3] = 4
逻辑视图:
[[1, 2, 3],
[4, 5, 6]]
─────────────────────────────────────
转置后 y (shape=(3,2), stride=(1,3)):
Storage: [1, 2, 3, 4, 5, 6] ← 完全相同的存储!
投影规则: y[i,j] = storage[i×1 + j×3]
y[0,0] = storage[0×1 + 0×3] = storage[0] = 1
y[0,1] = storage[0×1 + 1×3] = storage[3] = 4
y[1,0] = storage[1×1 + 0×3] = storage[1] = 2
逻辑视图:
[[1, 4],
[2, 5],
[3, 6]]
实现原理: 通过改变投影规则,同一个 Storage 被"重新解释"了!
x = torch.arange(12).view(3, 4)
# x: shape=(3, 4), stride=(4, 1), offset=0
# Storage: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
# 逻辑视图:
# tensor([[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
# 切片:取第 1 列(第 2 列,索引从 0 开始)
y = x[:, 1]
# 含义:取所有行的第 1 列
# 结果:[x[0,1], x[1,1], x[2,1]] = [1, 5, 9]
print(y) # tensor([1, 5, 9])
print("Shape:", y.shape) # (3,)
print("Stride:", y.stride()) # (4,) ← 每次跳 4 个元素
print("Storage offset:", y.storage_offset()) # 1 ← 从 storage[1] 开始
# 验证:共享 Storage(没有复制数据)
print(x.storage().data_ptr() == y.storage().data_ptr()) # True
# 验证:修改 y 会影响 x
y[0] = 999
print(x)
# tensor([[ 0, 999, 2, 3], ← x[0,1] 变成了 999
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
实现原理:
三要素的变化:
访问公式:
y[i] = storage[offset + i×stride[0]]
= storage[1 + i×4]
y[0] = storage[1 + 0×4] = storage[1] = 1
y[1] = storage[1 + 1×4] = storage[5] = 5
y[2] = storage[1 + 2×4] = storage[9] = 9
可视化:
Storage: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
↑ ↑ ↑
offset +4跳 +4跳
y[0]=1 y[1]=5 y[2]=9
x = torch.tensor([1, 2, 3]) # shape: (3,)
y = x.unsqueeze(0) # shape: (1, 3)
print("Shape:", y.shape) # (1, 3)
print("Stride:", y.stride()) # (3, 1)
实现原理:
x: shape=(3,), stride=(1,)
[1, 2, 3]
y: shape=(1, 3), stride=(3, 1)
[[1, 2, 3]]
y[0, 0] = storage[0×3 + 0×1] = storage[0] = 1
y[0, 1] = storage[0×3 + 1×1] = storage[1] = 2
x = torch.tensor([[1], [2], [3]])
# x: shape=(3, 1), stride=(1, 1), Storage=[1, 2, 3]
# 含义:3 行 1 列的矩阵
# [[1],
# [2],
# [3]]
y = x.expand(3, 4)
# y: shape=(3, 4), stride=(1, 0) ← 关键:第 1 维 stride=0!
# 含义:扩展成 3 行 4 列,但不复制数据
# Storage 还是 [1, 2, 3],只有 3 个元素!
print(y)
# tensor([[1, 1, 1, 1], ← 第 0 行重复 storage[0]
# [2, 2, 2, 2], ← 第 1 行重复 storage[1]
# [3, 3, 3, 3]]) ← 第 2 行重复 storage[2]
print("Stride:", y.stride()) # (1, 0) ← 第 1 维 stride 是 0!
print("共享 Storage:", x.storage().data_ptr() == y.storage().data_ptr()) # True
# 验证:修改 x 会影响 y 的整行
x[0, 0] = 999
print(y)
# tensor([[999, 999, 999, 999], ← 整行都变了!
# [ 2, 2, 2, 2],
# [ 3, 3, 3, 3]])
神奇之处:Stride = (1, 0) 的魔法
访问公式:
y[i, j] = storage[i×1 + j×0] = storage[i]
关键:j×0 = 0,无论 j 是多少!
第 0 行(i=0):
y[0, 0] = storage[0×1 + 0×0] = storage[0] = 1
y[0, 1] = storage[0×1 + 1×0] = storage[0] = 1 ← 和上面一样!
y[0, 2] = storage[0×1 + 2×0] = storage[0] = 1 ← 还是一样!
y[0, 3] = storage[0×1 + 3×0] = storage[0] = 1 ← 都访问 storage[0]
第 1 行(i=1):
y[1, 0] = storage[1×1 + 0×0] = storage[1] = 2
y[1, 1] = storage[1×1 + 1×0] = storage[1] = 2 ← 都访问 storage[1]
y[1, 2] = storage[1×1 + 2×0] = storage[1] = 2
y[1, 3] = storage[1×1 + 3×0] = storage[1] = 2
可视化:
Storage: [1, 2, 3] ← 只有 3 个元素
↓ ↓ ↓
逻辑视图:
j=0 j=1 j=2 j=3
i=0 1 1 1 1 ← 都指向 storage[0]
i=1 2 2 2 2 ← 都指向 storage[1]
i=2 3 3 3 3 ← 都指向 storage[2]
内存节省:
x = torch.randn(4, 4)
# 以下操作都是 View(零拷贝):
y1 = x.view(16) # 改变形状
y2 = x.reshape(2, 8) # 大多数情况是 View
y3 = x.transpose(0, 1) # 转置
y4 = x.permute(1, 0) # 维度重排
y5 = x[:, 2:] # 切片
y6 = x.unsqueeze(0) # 增加维度
y7 = x.squeeze() # 删除大小为 1 的维度
y8 = x.expand(4, 4, 4) # 虚拟扩展
y9 = x[::2, ::2] # 步进切片
# 验证:检查底层指针
print(x.data_ptr() == y1.data_ptr()) # True
x = torch.randn(4, 4)
# 以下操作会复制数据:
y1 = x.clone() # 显式复制
y2 = x.contiguous() # 如果不连续,则复制
y3 = x + 1 # 任何算术运算
y4 = x[x > 0] # 布尔索引
y5 = x[[0, 2]] # 高级索引(非连续选择)
y6 = torch.cat([x, x]) # 拼接
# 验证:不同的内存地址
print(x.data_ptr() == y1.data_ptr()) # False
小结: View 操作通过巧妙地改变投影规则(Stride、Offset),实现了零拷贝。这是 PyTorch 高性能的秘密。
但这引出了新问题:既然 transpose 只改变 Stride,为什么有时候需要调用 contiguous()?什么是"连续"?
这就是 Part 4 要解答的...
一个 Tensor 是 contiguous 的,当且仅当它的元素在 Storage 中是按照**标准顺序(C-order,行优先)**排列的。
标准顺序的定义:
# 对于 shape = (d₀, d₁, d₂, ..., dₙ)
# 标准 stride(C-order,行优先):
stride[n] = 1
stride[n-1] = d[n]
stride[n-2] = d[n-1] × d[n]
...
stride[0] = d[1] × d[2] × ... × d[n]
示例:
x = torch.tensor([[1, 2, 3],
[4, 5, 6]])
print(x.shape) # (2, 3)
print(x.stride()) # (3, 1)
# 检查是否连续:
expected_stride = (3, 1) # 3 = 第 1 维的大小, 1 = 最后一维
print(x.is_contiguous()) # True
直观理解:
Contiguous 意味着:
- 逻辑顺序:[[1,2,3], [4,5,6]] → 扁平化 → [1,2,3,4,5,6]
- Storage 顺序:[1, 2, 3, 4, 5, 6]
- 完美匹配!
x = torch.tensor([[1, 2, 3],
[4, 5, 6]])
y = x.transpose(0, 1)
print(y.is_contiguous()) # False
# 检查 y 的逻辑顺序:
print(y)
# tensor([[1, 4],
# [2, 5],
# [3, 6]])
# 逻辑顺序:[[1,4], [2,5], [3,6]] → 扁平化 → [1,4,2,5,3,6]
# Storage 顺序:[1, 2, 3, 4, 5, 6]
# 不匹配!
可视化:
Storage: [1, 2, 3, 4, 5, 6]
↓ ↓ ↓ ↓
逻辑读取: 1 → 4 2 → 5 ...
↑ ↑ ↑
跳跃访问!不连续!
案例:view() 要求输入连续
x = torch.randn(3, 4)
y = x.transpose(0, 1) # y 不连续
# 这会报错!
try:
z = y.view(12)
except RuntimeError as e:
print(e)
# RuntimeError: view size is not compatible with input tensor's
# size and stride. Use .reshape(...) or .contiguous().view(...)
为什么 view 要求连续?
view 只改变元数据,不移动数据。如果 Tensor 不连续,就无法通过简单的 stride 计算来重新解释 Storage。
例子:
y 的 Storage: [1, 2, 3, 4, 5, 6]
y 的逻辑顺序: [1, 4, 2, 5, 3, 6]
如果要 view 成 (12,):
需要访问顺序: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
但 Storage 只有 6 个元素!
无法用简单的 stride 公式实现 → 必须复制数据
contiguous() 做了什么?y = x.transpose(0, 1) # 不连续
z = y.contiguous() # 创建新的连续 Storage
print("y 的 Storage 指针:", y.data_ptr())
print("z 的 Storage 指针:", z.data_ptr())
# 不同!z 是新的 Tensor
print("z 是否连续:", z.is_contiguous()) # True
contiguous() 的工作流程:
# 源码简化版:
def contiguous(tensor):
if tensor.is_contiguous():
return tensor # 已经连续,不复制
# 创建新 Storage,按逻辑顺序复制数据
new_storage = allocate_storage(tensor.numel())
for idx in ndindex(tensor.shape): # 遍历所有索引
new_storage[flat_index(idx)] = tensor[idx]
return Tensor(new_storage, tensor.shape, standard_stride(tensor.shape))
方案 1:先 contiguous,再 view
y = x.transpose(0, 1) # 不连续
z = y.contiguous().view(12) # OK
方案 2:使用 reshape(推荐)
y = x.transpose(0, 1) # 不连续
z = y.reshape(12) # OK,自动处理
reshape vs view:
| 特性 | view | reshape |
|---|---|---|
| 要求连续 | 是 | 否 |
| 总是返回 View | 是 | 不一定 |
| 性能 | 最快 | 智能选择 |
| 推荐 | 性能关键路径 | 日常使用 |
import time
x = torch.randn(1000, 1000)
y = x.t() # 不连续
# 测试访问速度
start = time.perf_counter()
_ = x.sum()
time_cont = time.perf_counter() - start
start = time.perf_counter()
_ = y.sum()
time_non_cont = time.perf_counter() - start
print(f"连续 Tensor: {time_cont * 1000:.2f} ms")
print(f"不连续 Tensor: {time_non_cont * 1000:.2f} ms")
print(f"慢了: {time_non_cont / time_cont:.2f}x")
# 不连续的 Tensor 通常慢 1.5-3x(缓存不友好)
为什么慢?
小结: View 操作的代价是可能产生不连续的 Tensor,某些操作需要先调用 contiguous() 复制数据。
现在你理解了 Tensor 的底层机制。接下来,让我们系统学习 PyTorch 提供的各种维度操作工具。
view() - 严格的零拷贝x = torch.arange(12)
# view 要求 Tensor 连续
y = x.view(3, 4) # OK
y = x.view(2, 6) # OK
y = x.view(-1) # OK,-1 表示自动推断
# 元素总数必须匹配
try:
y = x.view(3, 5) # 错误!3×5=15 ≠ 12
except RuntimeError as e:
print(e)
reshape() - 智能的形状变换x = torch.randn(3, 4)
y = x.transpose(0, 1) # y 不连续
# reshape 会自动处理:
# - 如果可以,返回 View(零拷贝)
# - 如果不行,自动调用 contiguous()
z = y.reshape(12) # OK,自动处理了不连续的情况
flatten() - 扁平化x = torch.randn(2, 3, 4)
# 从某个维度开始展平
y = x.flatten(start_dim=1) # shape: (2, 12)
# 等价于:x.reshape(2, -1)
unsqueeze() - 添加大小为 1 的维度x = torch.randn(3, 4) # shape: (3, 4)
y = x.unsqueeze(0) # shape: (1, 3, 4)
z = x.unsqueeze(1) # shape: (3, 1, 4)
w = x.unsqueeze(-1) # shape: (3, 4, 1)
工程用途:
# 场景 1:匹配批次维度
image = torch.randn(3, 224, 224) # 单张图片 (C, H, W)
batch = image.unsqueeze(0) # (1, C, H, W)
# 场景 2:为广播做准备
a = torch.randn(5) # (5,)
b = torch.randn(3) # (3,)
a = a.unsqueeze(1) # (5, 1)
b = b.unsqueeze(0) # (1, 3)
c = a * b # (5, 3) - 广播乘法
squeeze() - 删除大小为 1 的维度x = torch.randn(1, 3, 1, 4)
y = x.squeeze() # (3, 4) - 删除所有大小为 1 的维度
z = x.squeeze(0) # (3, 1, 4) - 只删除第 0 维
w = x.squeeze(2) # (1, 3, 4) - 只删除第 2 维
transpose() - 交换两个维度x = torch.randn(2, 3, 4)
y = x.transpose(0, 1) # shape: (3, 2, 4)
z = x.transpose(1, 2) # shape: (2, 4, 3)
原理: 交换对应的 shape 和 stride
x: shape=(2, 3, 4), stride=(12, 4, 1)
y: shape=(3, 2, 4), stride=(4, 12, 1) # 交换了前两个
permute() - 任意维度重排x = torch.randn(2, 3, 4, 5)
# 重新排列所有维度
y = x.permute(3, 1, 0, 2) # shape: (5, 3, 2, 4)
# 常见用法:(B, C, H, W) → (B, H, W, C)
images = torch.randn(32, 3, 224, 224) # PyTorch 格式
images_hwc = images.permute(0, 2, 3, 1) # (32, 224, 224, 3)
Broadcasting 是 PyTorch 最强大的特性之一。
从右向左对齐维度,然后:
matrix = torch.tensor([[1, 2, 3],
[4, 5, 6]]) # (2, 3)
vector = torch.tensor([10, 20, 30]) # (3,)
result = matrix + vector # (2, 3)
# 对齐过程:
# matrix: (2, 3)
# vector: (3,) → 自动变为 (1, 3) → broadcast 到 (2, 3)
a = torch.randn(5, 1, 3) # (5, 1, 3)
b = torch.randn(1, 4, 3) # (1, 4, 3)
c = a + b # shape: (5, 4, 3)
# 对齐过程:
# a: (5, 1, 3)
# b: (1, 4, 3)
# ↓ ↓ ↓
# c: (5, 4, 3)
expand() - 显式 Broadcastingx = torch.tensor([[1], [2], [3]]) # shape: (3, 1)
# 虚拟扩展(不复制数据)
y = x.expand(3, 4) # shape: (3, 4)
print(y)
# tensor([[1, 1, 1, 1],
# [2, 2, 2, 2],
# [3, 3, 3, 3]])
# 验证:Stride 的第 1 维是 0
print(y.stride()) # (1, 0)
| 操作 | 改变的元数据 | 是否 View | 说明 |
|---|---|---|---|
view(shape) | shape, stride | 要求连续 | |
reshape(shape) | shape, stride | 通常 | 自动处理不连续 |
transpose(a, b) | shape, stride | 交换两维 | |
permute(dims) | shape, stride | 重排所有维度 | |
x[a:b] | offset, shape, stride | 切片 | |
x[::step] | stride | 步进 | |
unsqueeze(dim) | shape, stride | 插入维度 | |
squeeze() | shape, stride | 删除大小为 1 的维度 | |
expand(...) | shape | 虚拟扩展(stride 设为 0) | |
flatten() | shape, stride | 通常 | 展平 |
小结: 掌握这些维度操作,你就可以灵活地重塑 Tensor,为计算做准备。
现在你已经理解了 Tensor 的所有核心机制。接下来,让我们看看如何将 Tensor 组织成神经网络模型。
import torch
import torch.nn as nn
class SimpleModule(nn.Module):
def __init__(self, input_dim, output_dim):
super().__init__() # 必须调用!
# 可训练参数
self.weight = nn.Parameter(torch.randn(input_dim, output_dim))
self.bias = nn.Parameter(torch.zeros(output_dim))
# 不可训练的状态(如 BatchNorm 的 running_mean)
self.register_buffer('running_sum', torch.zeros(output_dim))
def forward(self, x):
output = torch.matmul(x, self.weight) + self.bias
self.running_sum += output.sum(dim=0)
return output
| 类型 | 用途 | 是否可训练 | 是否保存 | 是否移动到 device |
|---|---|---|---|---|
| Parameter | 模型参数 | |||
| Buffer | 模型状态 | |||
| 普通 Tensor | 临时变量 |
实例:
class Model(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.randn(10, 5))
self.register_buffer('mean', torch.zeros(5))
self.scale = torch.tensor(2.0)
model = Model()
# 移动到 GPU
model = model.cuda()
print(model.weight.device) # cuda:0
print(model.mean.device) # cuda:0
print(model.scale.device) # cpu ← 普通 Tensor 不会移动!
class MyModule(nn.Module):
def __init__(self):
super().__init__()
# 自动注册为子模块
self.layer1 = nn.Linear(10, 20)
self.layer2 = nn.Linear(20, 5)
model = MyModule()
# 自动递归收集所有参数
print(len(list(model.parameters()))) # 4 个(2 个 weight + 2 个 bias)
# 错误示例
class BadModule(nn.Module):
def __init__(self):
super().__init__()
self.layers = [nn.Linear(10, 10) for _ in range(3)]
model = BadModule()
print(len(list(model.parameters()))) # 0 ← 参数丢失!
# 正确示例
class GoodModule(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList([
nn.Linear(10, 10) for _ in range(3)
])
model = GoodModule()
print(len(list(model.parameters()))) # 6
小结: nn.Module 提供了一个优雅的容器,自动管理参数、状态和子模块。
现在,让我们通过实战练习巩固所有知识。
目标: 不使用 nn.Linear,只用 torch.matmul 实现两层 MLP。
import torch
import torch.nn as nn
class ManualMLP(nn.Module):
"""手动实现的多层感知机"""
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
# 第一层权重和偏置
self.w1 = nn.Parameter(
torch.randn(input_dim, hidden_dim) * (2.0 / input_dim) ** 0.5
)
self.b1 = nn.Parameter(torch.zeros(hidden_dim))
# 第二层权重和偏置
self.w2 = nn.Parameter(
torch.randn(hidden_dim, output_dim) * (2.0 / hidden_dim) ** 0.5
)
self.b2 = nn.Parameter(torch.zeros(output_dim))
def forward(self, x):
"""
Args:
x: shape (batch_size, input_dim)
Returns:
output: shape (batch_size, output_dim)
"""
# 第一层:x @ W1 + b1
hidden = torch.matmul(x, self.w1) + self.b1
# 激活函数:ReLU
hidden = torch.clamp(hidden, min=0)
# 第二层:hidden @ W2 + b2
output = torch.matmul(hidden, self.w2) + self.b2
return output
# 测试
model = ManualMLP(784, 128, 10)
x = torch.randn(32, 784) # 模拟 MNIST
output = model(x)
print(output.shape) # torch.Size([32, 10])
目标: 理解 Transformer 的核心,掌握复杂的维度变换。
import torch
import torch.nn as nn
import torch.nn.functional as F
class ManualMultiHeadAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
assert embed_dim % num_heads == 0
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
# Q、K、V 的投影矩阵
self.w_q = nn.Parameter(torch.randn(embed_dim, embed_dim) * 0.02)
self.w_k = nn.Parameter(torch.randn(embed_dim, embed_dim) * 0.02)
self.w_v = nn.Parameter(torch.randn(embed_dim, embed_dim) * 0.02)
self.w_o = nn.Parameter(torch.randn(embed_dim, embed_dim) * 0.02)
def forward(self, x):
"""
Args:
x: (batch, seq_len, embed_dim)
Returns:
output: (batch, seq_len, embed_dim)
"""
batch_size, seq_len, embed_dim = x.shape
# 步骤 1:线性投影
Q = torch.matmul(x, self.w_q) # (B, L, E)
K = torch.matmul(x, self.w_k)
V = torch.matmul(x, self.w_v)
# 步骤 2:切分成多头
Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim)
Q = Q.transpose(1, 2) # (B, H, L, D)
K = K.view(batch_size, seq_len, self.num_heads, self.head_dim)
K = K.transpose(1, 2)
V = V.view(batch_size, seq_len, self.num_heads, self.head_dim)
V = V.transpose(1, 2)
# 步骤 3:计算注意力分数
scores = torch.matmul(Q, K.transpose(-2, -1)) # (B, H, L, L)
scores = scores / (self.head_dim ** 0.5)
# 步骤 4:Softmax
attn_weights = F.softmax(scores, dim=-1)
# 步骤 5:加权求和
attn_output = torch.matmul(attn_weights, V) # (B, H, L, D)
# 步骤 6:合并多头
attn_output = attn_output.transpose(1, 2) # (B, L, H, D)
attn_output = attn_output.contiguous()
attn_output = attn_output.view(batch_size, seq_len, embed_dim)
# 步骤 7:输出投影
output = torch.matmul(attn_output, self.w_o)
return output
# 测试
model = ManualMultiHeadAttention(embed_dim=512, num_heads=8)
x = torch.randn(2, 10, 512)
output = model(x)
print(output.shape) # torch.Size([2, 10, 512])
关键维度变换可视化:
输入: (B, L, E)
↓ Linear projection
Q, K, V: (B, L, E)
↓ view + transpose
Q, K, V: (B, H, L, D) 其中 H=num_heads, D=E/H
↓ Attention 计算
Output: (B, H, L, D)
↓ transpose + contiguous + view
Output: (B, L, E)
def manual_layer_norm(x, eps=1e-5):
"""
Args:
x: shape (batch, ..., features)
Returns:
normalized: shape (batch, ..., features)
"""
# 在最后一维计算统计量
mean = x.mean(dim=-1, keepdim=True) # (B, ..., 1)
var = x.var(dim=-1, keepdim=True, unbiased=False)
# 归一化
x_normalized = (x - mean) / torch.sqrt(var + eps)
return x_normalized
# 测试
x = torch.randn(32, 10, 512)
output = manual_layer_norm(x)
# 验证
print(output[0, 0].mean()) # 接近 0
print(output[0, 0].std()) # 接近 1
Tensor = Storage(物理)+ (Shape, Stride, Offset)(投影)
T[i,j,k] = S[offset + i×stride[0] + j×stride[1] + k×stride[2]]分离式设计的动机
View 操作的原理
Contiguous 的必要性
维度操作的统一理解
nn.Module 的设计
reshape() 而非 view()ModuleList 管理子模块transpose 后可能需要 contiguous()在下一篇文章中,我们将深入自动微分引擎(Autograd),理解:
backward() 的实现原理在进入下一阶段前,请确保你能回答:
如果你能清晰回答这些问题,恭喜你,你已经掌握了 PyTorch 的核心机制!
下一篇预告: 《阶段二:后向传播 - Autograd 自动微分引擎揭秘》
本文是 PyTorch 架构级学习系列的第 2 篇。