expand与repeat:省复制,为何梯度仍相加
只把大小为1的维度读成更大形状,优先考虑expand;需要真实复制数据,或准备给每个副本单独写值,再考虑repeat或显式复制。两者即使打印出完全相同的数组,也可能拥有不同的存储关系。更容易误判的是反向:复制数据不等于创建彼此独立的模型参数。
文章目录
下面用一个1×3向量扩成4×3,分别看地址、写入、梯度和下游内存。目标是回答“expand和repeat怎么选”,而不是给所有训练任务推荐同一个更快写法。Read in English。
四行数字,可能只对应三个存储位置
PyTorch 2.8的expand文档把它定义为视图:沿大小为1的轴扩展时,用零步长复用已有存储。对x=[[1,2,3]],把形状变成(4,3),逻辑上出现四行,但每行仍指向同一组三个数。
以连续float64源数组为例,原来一行中相邻元素相隔8字节。扩展后的地址偏移可写成i×0+j×8;行号i变化不挪动地址,列号j才决定读哪个元素。这里的地址偏移以字节计;NumPy的strides也是字节,而PyTorch的stride通常以元素个数计。同一布局可分别记作(0,8)字节和(0,1)元素,日志里比较前先统一单位。
它不是先分配十二个数,再压缩成三个,而是从一开始就用同一个位置回答多次读取。
这也解释了适用条件:原有大小为2的轴,不能用expand直接扩成4;两个不同元素的意义不能被“反复读同一个位置”代替。显式重复可以另行定义,但那是另一种操作。先写出源形状与目标形状,再从尾部对齐检查每个维度,比根据打印结果猜行为可靠。
repeat复制数据,不会自动创造独立参数
repeat文档明确说它复制数据,并提示它更接近NumPy的tile,而非名字相似的numpy.repeat。本文用tile构造同值、不同存储的对照,不把两个库的同名函数当成相同接口。
设两条程序都从同一个变量x产生四份数值,再计算L=Σᵢⱼwᵢⱼyᵢⱼ。无论中间是共享视图还是按原值复制,作为x的函数都有yᵢⱼ=xⱼ,因此∂L/∂xⱼ=Σᵢwᵢⱼ。取权重四行依次为(1,2,3)、(4,5,6)、(7,8,9)、(10,11,12),对x的导数是(22,26,30)。
求导时问的是同一个输入变量改变,会影响多少处使用;存储时问的是这些值目前放在哪里。分配了新内存,并不会让输入变量的依赖关系消失。不要因为看到repeat产生副本,就把它理解成四个能各自学习的参数。若模型确实需要独立参数,应在模型参数化中明确表达,而不是从复制操作推断。
这里的导数由链式法则推导,并用有限差分比较两条数值程序;没有运行PyTorch自动微分。对PyTorch具体算子是否可导、叶子张量和优化器注册等问题,需要另查相应语义。站内detach与clone的区别进一步区分存储关系与梯度边界。
共享视图适合读,逐副本修改先分离存储
如果四行都引用同一位置,“只把第0行第0列改成9”可能无法表达为一个独立副本的更新。PyTorch的expand文档因此警告重叠存储上的原地操作,并建议需要写入时先clone。不要依赖某个向量化原地操作刚好报错或刚好得到预期结果。
本例用NumPy的broadcast_to执行地址实验。它返回只读广播视图,这是NumPy接口的保护策略,不能写成“PyTorch的expand也是只读”。更清楚的验证方式是修改源数组:视图的四行一起读到新值,而此前tile得到的副本仍保留旧值。
真正需要逐行改值时,先显式复制,再修改其中一个位置。代码中c=y.copy()后改c[0,0],不会改变源数组,也不会改变c的其他行。这证明这一次复制后的存储独立;不据此推断自动微分被切断。关于是否保留求导路径,需与存储是否分离分别判断。
扩展不复制,不代表整条算子链不分配
4×3的float64逻辑数组有12个元素,按元素大小算是96字节;源数组只有24字节。广播视图的nbytes仍会报告逻辑元素总量,不能拿这个值断言它另行占用了96字节数据,也不能拿源数组24字节断言整个表达式只用24字节。
代码中的out=y+ones生成新的4×3输出,实际有96字节元素数据,并且不再与x共享。类似地,扩展后接一个必须形成稠密结果的运算,省下的可能只是输入复制;下游输出、临时缓冲和反向保存量仍需分别核算。
做一个明确标注为估算的预算:形状(B,T,D)=(32,2048,4096),若元素为2字节,完整数据量是536870912字节,即512MiB;一个(1,1,D)源向量是8192字节。扩展输入可以避免先复制出那512MiB,但如果下游必须产生该形状的稠密输出,输出预算仍在。这个计算不是GPU分配器实测,也没有包含其他张量、缓存和开销。
若接着reshape,不能靠函数名保证仍是零复制;形状与步长能否兼容很重要。已有view与reshape的存储判断讨论这一层。这里重点是复用同一输入,并没有重新讲一遍所有形状变换。
把存储、值和导数分开验证
下面的程序实际运行NumPy数组:检查形状与数值一致、零行步长、源与视图共享、源与副本不共享、只读标记、源修改传播、复制后单独写入、逻辑字节数、下游输出分配、不合法扩展,以及两条路径各自的有限差分。十四项断言都服务于上面的区分。
# NumPy storage experiment and explicit derivatives, not PyTorch autograd.
import numpy as np
x=np.array([[1.,2.,3.]])
y=np.broadcast_to(x,(4,3))
r=np.tile(x,(4,1))
assert y.shape==r.shape==(4,3)
assert np.array_equal(y,r)
assert y.strides[0]==0
assert np.shares_memory(x,y)
assert not np.shares_memory(x,r)
assert not y.flags.writeable # NumPy policy; not a PyTorch guarantee.
x[0,0]=9.
assert np.all(y[:,0]==9.) and np.all(r[:,0]==1.)
c=y.copy(); c[0,0]=-1.
assert x[0,0]==9. and c[1,0]==9.
assert y.nbytes==r.nbytes==96 and x.nbytes==24
out=y+np.ones((4,3))
assert out.nbytes==96 and not np.shares_memory(x,out)
try:
np.broadcast_to(np.ones((2,3)),(4,3))
except ValueError:
pass
else:
raise AssertionError('A non-singleton axis cannot expand this way')
x=np.array([[1.,2.,3.]])
w=np.array([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.],[10.,11.,12.]])
expected=w.sum(axis=0,keepdims=True)
assert expected.tolist()==[[22.,26.,30.]]
def loss(z,kind):
a=np.broadcast_to(z,(4,3)) if kind=='view' else np.tile(z,(4,1))
return float((a*w).sum())
for kind in ('view','copy'):
numerical=np.empty_like(x)
for j in range(3):
a=x.copy(); b=x.copy(); a[0,j]+=1e-5; b[0,j]-=1e-5
numerical[0,j]=(loss(a,kind)-loss(b,kind))/(2e-5)
assert np.allclose(numerical,expected,rtol=1e-8,atol=1e-8)
print('14 storage and derivative checks passed; NumPy',np.__version__)
print('source bytes',x.nbytes,'logical broadcast bytes',y.nbytes)
有限差分每次只扰动x的一列,权重和目标形状保持不变。两条程序的标量目标相同,数值导数应与手工求和(22,26,30)一致。当前环境未安装PyTorch;这不是Tensor.expand、Tensor.repeat、backward或CUDA性能的执行记录。
比较前还要固定归约方式:若把整个4×3加权和改成所有12项的平均,导数也要除以12,变成(11/6,13/6,5/2)。这个变化来自目标函数,不是expand“导致梯度变大”。重复使用次数、权重、sum或mean都应出现在推导里。
从使用需求决定,不从API名字决定
只读地把一份偏置或共享特征用于多个位置,先确认广播轴确实为1,再看下游是否接受该布局;很多支持广播的算子甚至不需要提前显式expand。需要输出中各份数值能单独修改,就为独立存储留预算。需要每份独立学习,还要重新审视模型中参数是否共享。
实际排查时记录三组证据:输入输出形状和步长;源、副本、下游输出之间的共享关系;固定损失下的梯度或有限差分。性能问题再测完整前向与反向路径、实际设备和峰值内存。本文没有给出“expand必然更快”或“repeat必然更慢”的结论。
如果这组地址推导和十四项检查帮你分清了省复制、独立写入与共享梯度,欢迎在文末自愿支持。例子可以直接缩成你自己的输入形状,先验证含义,再比较性能。
来源与验证范围
PyTorch 2.8 expand · PyTorch 2.8 repeat · PyTorch broadcasting · PyTorch reshape · NumPy 2.3 broadcast_to
依据PyTorch 2.8版本文档核对,广播说明更新于2021年1月31日;这是一篇工程原理文章,不作为新发布新闻。原创例子实际运行NumPy,未运行PyTorch或设备基准。计划文章日期为2026年10月22日23:30北京时间,排期与实际公开另行核验。
支持
如果文中的解释或示例帮你解决了问题、节省了时间,欢迎自愿打赏支持本站。
二维码可点击放大。更多支持方式见支持页面。


