在 PyTorch 中有条件地应用张量操作

Question

我知道 PyTorch 没有类似 map 的函数来将函数应用于张量的每个元素。那么，如果 PyTorch 中没有类似 map 的函数，我可以执行以下操作吗？

if tensor_a * tensor_b.matmul(tensor_c) < 1:
    return -tensor_a*tensor_b
else:
    return 0

如果张量是一维的，这将起作用。但是，当tensor_b 是2D（tensor_a 需要在return 语句中为unsqueezed）时，我需要它来工作。这意味着应该返回一个二维张量，其中一些行将是0 向量。

乐于使用最新 Python 版本的最新功能。

Answer 1

如果我理解正确，您希望以任何一种方式（因此映射）返回张量，但通过逐个检查条件来返回。假设tensor_a、tensor_b 和tensor_c 的形状都是二维的，如在“简单矩阵”中，这是一个可能的解决方案。

您正在寻找的可能是torch.where，它非常接近于基于条件的映射，它将返回一个值或另一个元素方式。

它的工作方式类似于torch.where(condition, value_if, value_else)，其中所有三个张量具有相同的形状（value_if 和value_else 实际上可以是浮点数，它们将被转换为张量，填充相同的值）。此外，condition 是一个 bool 张量，它定义了分配给输出张量的值：它是一个布尔掩码。

出于本示例的目的，我使用了随机张量：

>>> a = torch.rand(2, 2, dtype=float)*100
>>> b = torch.rand(2, 2, dtype=float)*0.01
>>> c = torch.rand(2, 2, dtype=float)*10

>>> torch.where(a*(b@c) < 1, -a*b, 0.)
tensor([[ 0.0000,  0.0000],
        [ 0.0000, -0.0183]], dtype=torch.float64)

更一般地说，如果tensor_a 和tensor_b 的形状为(m, n)，并且tensor_c 的形状为(n, m)，则由于操作限制，这将起作用。在你的实验中，我猜你只有列。