PyTorch Linear层为何以(out,in)存储权重并在前向传播中转置?
PyTorch Linear层权重存储形状的疑问
PyTorch的Linear层将weight属性存储为(out_features, in_features)形状,它的前向传播逻辑如下:
def forward(self, input: Tensor) -> Tensor: return F.linear(input, self.weight, self.bias)
这一实现等价于:
def forward(self, input: Tensor) -> Tensor: return torch.addmm(self.bias, input, self.weight.t())
这里存在一个疑问:为什么不把weight属性直接存储为(in_features, out_features)形状,这样就能省去前向和反向传播中的转置操作?
我针对两种存储方式做了基准测试——前向与反向传播各执行10000次迭代,结果发现预转置权重的版本每次迭代要快约10µs。测试用的是含100个特征、128行数据的随机张量,为保证公平性,两个Linear层配置了完全相同的参数(包括权重和偏置),并且作用于同一组随机数据。
内容的提问来源于stack exchange,提问作者Chris
相关产品推荐
相关产品推荐

