PyTorch中使用torch.diff时如何保持张量原始维度？

原张量维度为(3,3,3)，使用torch.diff计算差分后维度变为(3,2,3)，可以通过设置torch.diff的prepend参数，传入对应维度的0张量来补全开头，从而保持原始维度。

具体实现代码 #

import torch

# 初始化原张量
a = torch.rand(size=(3,3,3), dtype=torch.float32)

# 构造用于开头补0的张量，dim=1维度长度为1，其余维度与原张量一致
prepend_zero = torch.zeros(size=(3,1,3), dtype=torch.float32)

# 传入prepend参数计算差分
a_diff = torch.diff(a, n=1, dim=1, prepend=prepend_zero)

print(a_diff.shape)  # 输出: torch.Size([3, 3, 3])

更简洁的写法 #

可以利用原张量的切片来构造补0张量，避免手动指定维度：

import torch

a = torch.rand(size=(3,3,3), dtype=torch.float32)
# 取原张量dim=1维度的第一个切片形状，生成全0张量
prepend_zero = torch.zeros_like(a[:, :1, :])
a_diff = torch.diff(a, n=1, dim=1, prepend=prepend_zero)

print(a_diff.shape)  # 输出: torch.Size([3, 3, 3])

原理说明：prepend参数会将传入的张量拼接在原张量的指定维度开头，差分计算时基于拼接后的张量进行，最终结果的维度就会和原张量保持一致。

内容的提问来源于stack exchange，提问作者freak11