第一个反直觉:Tensor 不持有数据。 Tensor 是一组元数据(shape / stride / offset / dtype / device), 数据本身存在它指向的 Storage 里。 transpose、view、slice 之所以是 O(1) 操作, 就是因为它们只改 Tensor 元数据,不动 Storage 一个字节。

三件事:shape / stride / offset

Python 里印证一下

import torch

a = torch.arange(12, dtype=torch.float32).view(3, 4)
print(a.shape, a.stride(), a.storage_offset())
# torch.Size([3, 4])  (4, 1)  0   ← 行优先 row-major

b = a.T   # 转置:仅交换 stride
print(b.shape, b.stride(), b.is_contiguous())
# torch.Size([4, 3])  (1, 4)  False

# A 和 B 共享 Storage —— 验证
print(a.data_ptr() == b.data_ptr())   # True
print(a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr())  # True

c = a[1:, :2]   # slice:offset 变了,stride 不变
print(c.shape, c.stride(), c.storage_offset())
# torch.Size([2, 2])  (4, 1)  4

# contiguous 会真的复制
d = b.contiguous()
print(a.data_ptr() == d.data_ptr())   # False —— 新 Storage

PyTorch 的一个 op(比如 torch.add)其实有几十个不同的实现—— CPU 一个、CUDA 一个、Vulkan 一个、Autograd 包装一个、Autocast 包装一个…… Dispatcher 就是根据当前 Tensor 的属性挑选哪一个实现的中间路由层。 它是 PyTorch 全部扩展性的来源。

常见 DispatchKey 速查

native_functions.yaml — 算子签名的真理之源

所有 ATen 算子的签名、变体、dispatch 配置都集中在 aten/src/ATen/native/native_functions.yaml。读 PyTorch 源码的入口就在这里。

# aten/src/ATen/native/native_functions.yaml 节选 (示意)
- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  device_check: NoCheck
  structured_delegate: add.out
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA: add_sparse
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
  tags: pointwise

# 说明:
#   structured_delegate: 这个 op 是 structured kernel,实际实现复用 add.out
#   dispatch:            列出每种 DispatchKey 对应的 kernel 函数名
#   未列出的 backend     由代码生成器自动生成默认派发

Autograd 是 PyTorch 最迷人的部分——一行 loss.backward() 让所有参数自动有了梯度。 今天只做机制概览(Day 09 才手写), 重点理解三个概念:requires_grad、grad_fn、动态计算图。

最小可演示

import torch

a = torch.tensor([2.0], requires_grad=True)
b = torch.tensor([3.0], requires_grad=True)
c = torch.tensor([4.0], requires_grad=True)

z = (a + b) * c

# 看一下 z 的 grad_fn —— 反向图的入口
print(z.grad_fn)
# <MulBackward0 object at 0x...>

print(z.grad_fn.next_functions)
# ((<AddBackward0 object>, 0), (<AccumulateGrad object>, 0))

# 触发反向
z.backward()

# 梯度被累加到叶子上
print(a.grad, b.grad, c.grad)
# tensor([4.])  tensor([4.])  tensor([5.])
# ∂z/∂a = c = 4, ∂z/∂b = c = 4, ∂z/∂c = a+b = 5

把前三节串起来。一行 c = a + b(a, b 是 CUDA Tensor,requires_grad=True), Python 解释器到 GPU 的 ALU 之间至少经过 8 步。这是 PyTorch 内部机制的体检图。

理论懂了没用——必须亲手看一遍 Dispatcher 在跑什么。两个实验: ① 用 __torch_dispatch__ 拦截所有 op 调用,把整个网络的 op 序列打印出来; ② 写一个 forward hook,在 ResNet 的某层加观察点。

实验 1 — TorchDispatchMode 打印所有 op

import torch
from torch.utils._python_dispatch import TorchDispatchMode

class OpPrinter(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        # func 是被调用的 ATen op,比如 aten.add.Tensor
        kwargs = kwargs or {}
        in_shapes = [tuple(a.shape) if isinstance(a, torch.Tensor) else a
                     for a in args]
        out = func(*args, **kwargs)
        out_shape = out.shape if isinstance(out, torch.Tensor) else "-"
        print(f"{func.__name__:<30} in={in_shapes}  out={out_shape}")
        return out

# 用它包住任意计算,所有 op 的真实调用都会被打印
with OpPrinter():
    a = torch.randn(4, 8, device="cuda")
    b = torch.randn(4, 8, device="cuda")
    c = (a + b).relu().sum()

# 你会看到 aten.add.Tensor、aten.relu、aten.sum.default 全部被记录
# 这就是 PyTorch 在 dispatcher 上真正执行的算子序列

实验 2 — 在真实模型上拦截

import torchvision

model = torchvision.models.resnet18().cuda().eval()
x = torch.randn(1, 3, 224, 224, device="cuda")

from collections import Counter
op_count = Counter()

class OpCounter(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        op_count[func.__name__] += 1
        return func(*args, **(kwargs or {}))

with torch.no_grad(), OpCounter():
    _ = model(x)

for op, n in op_count.most_common(10):
    print(f"{n:>4d}  {op}")

# 典型输出:
#   20  convolution.default
#   17  batch_norm.default
#   17  relu.default
#    9  add.Tensor          ← residual connection!
#    1  adaptive_avg_pool2d.default
#    1  addmm.default

实验 3 — 验证 view 不复制 Storage

import torch

a = torch.zeros(3, 4, device="cuda")
print("a base ptr =", a.untyped_storage().data_ptr())

# 各种 view 操作
b = a.T                              # transpose
c = a.view(2, 6)                     # reshape
d = a[1:]                            # slice
e = a.expand(5, 3, 4)                  # broadcast

for name, t in [("b", b), ("c", c), ("d", d), ("e", e)]:
    same = t.untyped_storage().data_ptr() == a.untyped_storage().data_ptr()
    print(f"{name}: same storage = {same}, stride = {t.stride()}")

# 全部 True,只有 stride 不同 —— 验证 view = O(1) 元数据操作

# contiguous() 才会真复制
f = b.contiguous()
print("f: same storage =", f.untyped_storage().data_ptr() == a.untyped_storage().data_ptr())
# False —— 新 Storage 被分配

1. 用一张图说明 Tensor、Storage、shape、stride、offset 五个概念的关系,并写出 T[i,j,k] 的物理地址公式。
2. 给定 a = torch.arange(24).view(2,3,4),手算 a.transpose(0,2) 之后的 shape 和 stride,并验证。
3. 解释 DispatchKey 的"洋葱"结构:为什么 Autograd 不是独立系统,而是 Dispatcher 上的一组 key?这给扩展性带来什么好处?
4. 把 c = a + b 的 8 步调用链复述一遍,每一步写出对应的源码文件(粗略路径即可)。
5. 用 TorchDispatchMode 抓取一个 Linear 层的算子序列(y = linear(x)),解释为什么会看到 addmm 而不是 matmul。