torch具身常用算子

PyTorch核心张量算子

覆盖：扩散模型 / DiT / 图像视频 / 机器人动作序列

约定记号：

X,Y\boldsymbol{X},\boldsymbol{Y}X,Y 代表张量；N,B,C,T,H,WN,B,C,T,H,WN,B,C,T,H,W = 批次、通道、时序、高、宽

一、维度变换算子（特征排布，最高频）

1. tensor.view(*shape) / torch.reshape

数学：重塑张量元素排布，元素顺序不变

设原张量 X∈Ra×b×c\boldsymbol{X} \in \mathbb{R}^{a\times b\times c}X∈Ra×b×c

reshape(X, a⋅b, c)∈Rab×c \mathrm{reshape}(\boldsymbol{X},\,a\\cdot b,\\,c) \in \mathbb{R}^{ab\times c} reshape(X,a⋅b,c)∈Rab×c

限制：元素总数不变；view要求内存连续，reshape自动处理连续性。

场景：DiT时空维度打平、latent变形、动作序列维度重组。

2. torch.permute(d₀,d₁,...)

维度重排列，无复制（非连续张量）

例：X∈RB,C,H,W\boldsymbol{X}\in\mathbb{R}^{B,C,H,W}X∈RB,C,H,W

permute(X, B,H,W,C) \mathrm{permute}(\boldsymbol{X},\,B,H,W,C) permute(X,B,H,W,C)

场景：图像 BCHW ↔ BHWC、视频调换时序/空间维度、注意力维度对齐。

3. torch.unsqueeze(dim) / squeeze(dim)

unsqueeze(X, d):RD0×...Dk→RD0×⋯×1×...Dk \mathrm{unsqueeze}(\boldsymbol{X},\,d): \mathbb{R}^{D_0\times\dots D_k} \rightarrow \mathbb{R}^{D_0\times\dots\times 1\times\dots D_k} unsqueeze(X,d):RD0×...Dk→RD0×⋯×1×...Dk

新增长度为1的维度；squeeze移除长度=1维度。

场景：单张图像扩batch维、给latent增加时序维度适配视频模型输入。

4. torch.stack([X₁,X₂,...,Xₙ], dim)

新增维度堆叠

Xi∈RH,W\boldsymbol{X}_i\in\mathbb{R}^{H,W}Xi∈RH,W

Z=stack(X1,X2, 0)  ⟹  Z∈R2×H×W \boldsymbol{Z} = \mathrm{stack}(\\boldsymbol{X}_1,\\boldsymbol{X}_2,\,0) \implies \boldsymbol{Z}\in\mathbb{R}^{2\times H\times W} Z=stack(X1,X2,0)⟹Z∈R2×H×W

公式表达：

Z0=X1,Z1=X2 \boldsymbol{Z}0=\boldsymbol{X}_1,\quad \boldsymbol{Z}1=\boldsymbol{X}_2 Z0=X1,Z1=X2

👉 KV Cache组装、多帧latent构建时序维度。

5. torch.cat([X₁,X₂], dim)

现有维度拼接，不新增维度

X1,X2∈RB,C,H,W\boldsymbol{X}_1,\boldsymbol{X}_2\in\mathbb{R}^{B,C,H,W}X1,X2∈RB,C,H,W

Z=cat(X1,X2, 2)∈RB,C,2H,W \boldsymbol{Z} = \mathrm{cat}(\\boldsymbol{X}_1,\\boldsymbol{X}_2,\,2) \in \mathbb{R}^{B,C,2H,W} Z=cat(X1,X2,2)∈RB,C,2H,W

场景：拼接条件嵌入、前后帧特征、历史动作序列。

二、采样 & 扩散核心数值算子

1. torch.randn_like(X)

ε∼N(0,I) \boldsymbol{\varepsilon} \sim \mathcal{N}(0,\boldsymbol{I}) ε∼N(0,I)

标准高斯噪声，扩散模型根基：

xt=αˉtx0+1−αˉtε \boldsymbol{x}_t = \sqrt{\bar{\alpha}_t}\boldsymbol{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\varepsilon} xt=αˉt x0+1−αˉt ε

2. torch.lerp(a, b, weight) 线性插值

lerp(a,b,w)=(1−w)⋅a+w⋅b \mathrm{lerp}(\boldsymbol{a},\boldsymbol{b},w) = (1-w)\cdot \boldsymbol{a} + w\cdot \boldsymbol{b} lerp(a,b,w)=(1−w)⋅a+w⋅b

应用：扩散采样更新latent、流匹配积分、动作轨迹平滑、帧插值。

3. torch.clamp(x, min, max)

clamp(x,a,b)={ax<axa≤x≤bbx>b \mathrm{clamp}(x,a,b)= \begin{cases} a & x<a\\ x & a\le x\le b\\ b & x>b \end{cases} clamp(x,a,b)=⎩ ⎨ ⎧axbx<aa≤x≤bx>b

场景：VAE解码值域限制防止数值溢出黑图、机器人动作限幅。

4. torch.sqrt,exp⁡,log⁡\texttt{torch.sqrt},\exp,\logtorch.sqrt,exp,log

扩散调度器通用：

σt=1−αˉt,αˉt=∏s=1tαs \sigma_t=\sqrt{1-\bar{\alpha}_t},\quad \bar{\alpha}t=\prod{s=1}^t \alpha_s σt=1−αˉt ,αˉt=s=1∏tαs

Flow Matching / Cosmos世界模型大量使用对数、指数计算SNR、时间条件。

三、掩码与条件选择

1. torch.masked_fill(mask, value)

Yi,j={vmaski,j=TrueXi,jmaski,j=False \boldsymbol{Y}{i,j}= \begin{cases} v & \mathrm{mask}{i,j}=True\\ \boldsymbol{X}{i,j} & \mathrm{mask}{i,j}=False \end{cases} Yi,j={vXi,jmaski,j=Truemaski,j=False

场景：Inpaint区域填充噪声；注意力padding位置填充 −∞-\infty−∞。

2. torch.where(cond, x, y)

where(c,x,y)i={xici=Trueyici=False \mathrm{where}(c,\boldsymbol{x},\boldsymbol{y})_{i}= \begin{cases} \boldsymbol{x}_i & c_i=\mathrm{True}\\ \boldsymbol{y}_i & c_i=\mathrm{False} \end{cases} where(c,x,y)i={xiyici=Trueci=False

场景：根据掩码混合原图latent与噪声latent。

四、注意力矩阵运算（DiT / Transformer / KV Cache）

1. torch.matmul(a,b) / @

矩阵乘法

C=AB,Cij=∑kAikBkj \boldsymbol{C} = \boldsymbol{A}\boldsymbol{B},\quad C_{ij}=\sum_k A_{ik}B_{kj} C=AB,Cij=k∑AikBkj

自注意力核心：

Attn(Q,K,V)=softmax(QK⊤dh)V \mathrm{Attn}(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V})=\mathrm{softmax}\left(\frac{\boldsymbol{Q}\boldsymbol{K}^\top}{\sqrt{d_h}}\right)\boldsymbol{V} Attn(Q,K,V)=softmax(dh QK⊤)V

2. torch.transpose(d1,d2)

(X⊤)ij=Xji (\boldsymbol{X}^\top){ij}=X{ji} (X⊤)ij=Xji

常用写法：K.transpose(-1,-2) 适配QK相乘。

3. F.softmax(x, dim)

softmax(x)i=exi∑jexj \mathrm{softmax}(\boldsymbol{x})_i = \frac{e^{x_i}}{\sum_j e^{x_j}} softmax(x)i=∑jexjexi

五、重采样算子（图像/视频latent缩放）

F.interpolate

连续空间重采样，以bilinear为例：

f^(u,v)=∑i∑jwij⋅f(xi,yj) \hat{f}(u,v)=\sum_i\sum_j w_{ij}\cdot f(x_i,y_j) f^(u,v)=i∑j∑wij⋅f(xi,yj)

场景：VAE上下采样、latent分辨率调整、视频时序缩放。

六、时序序列算子（机器人动作 / 视频）

1. torch.diff(X, dim)

相邻差分

Δxt=xt+1−xt \Delta \boldsymbol{x}t = \boldsymbol{x}{t+1}-\boldsymbol{x}_t Δxt=xt+1−xt

用途：由位置序列求解速度、加速度。

2. torch.cumsum(X, dim)

累积和

cumsum(X)t=∑k=1tXk \mathrm{cumsum}(\boldsymbol{X})t=\sum{k=1}^t \boldsymbol{X}_k cumsum(X)t=k=1∑tXk

Flow Matching积分、由速度积分得到轨迹。

七、归一化变换（图像VAE预处理）

x.sub_(mean).div_(std)

x^=x−μσ \hat{x} = \frac{x-\mu}{\sigma} x^=σx−μ

反归一化：

x=x^⋅σ+μ x = \hat{x}\cdot\sigma + \mu x=x^⋅σ+μ

常用：图像 0,1↔−1,10,1\leftrightarrow-1,10,1↔−1,1 转换输入VAE。

八、内存连续性

tensor.contiguous()

张量底层内存由非连续 → 连续排布 。

数学上不改变张量数值，改变存储布局；

permute/transpose后直接view会报错，必须调用。

熟练运用PyTorch张量算子完成扩散模型、视频时序与机器人动作序列开发：

通过stack/cat/permute/view/unsqueeze实现张量维度重组；依托randn_like\mathrm{randn\_like}randn_like生成高斯噪声，结合lerp\mathrm{lerp}lerp线性插值实现扩散与流匹配采样迭代；利用clamp\mathrm{clamp}clamp控制数值范围，masked_fill/where\mathrm{masked\_fill/where}masked_fill/where实现条件生成与图像修复；基于matmul/transpose/softmax\mathrm{matmul/transpose/softmax}matmul/transpose/softmax搭建自注意力计算；使用interpolate\mathrm{interpolate}interpolate完成潜空间分辨率重采样；借助diff/cumsum\mathrm{diff/cumsum}diff/cumsum对动作序列做差分与积分运算；同时掌握contiguous()处理维度置换后的内存连续性问题。