Pytorch 对沿轴的张量列表求和

Question

我有一个相同形状的张量列表。我想沿轴对整个张量列表求和。是否torch.cumsum沿暗淡执行此操作？如果是这样，它需要将列表转换为单个张量并求和？

Answer 1

您不需要cumsum，sum是您的朋友，是的，您应该首先根据您的需要stack或cat根据您的需要将它们转换为单个张量，如下所示：

import torch
my_list = [torch.randn(3, 5), torch.randn(3, 5)]
result = torch.stack(my_list, dim=0).sum(dim=0).sum(dim=0)
print(result.shape) #torch.Size([5])