pytorch 中是否有元素映射 function？

Question

我是 PyTorch 的新手，我来自函数式编程语言（map function 到处都在使用）。 问题是我有一个张量，我想对张量的每个元素做一些操作。 操作可能多种多样，所以我需要这样的 function：

map : (Numeric -> Numeric) -> Tensor -> Tensor

e.g. map(lambda x: x if x < 255 else -1, tensor) # the example is simple, but the lambda may be very complex

PyTorch中有这样的function吗？ 我应该如何实现这样的 function？

Answer 1

为张量实现的大多数数学运算（对于 numpy 中的 ndarray 也是如此）实际上是按元素应用的，因此您可以编写例如

mask = tensor < 255
result = tensor * mask + (-1) * ~mask

这是一个相当普遍的方法。 对于您现在只想修改某些元素的情况，您还可以应用“逻辑索引”来覆盖当前张量：

tensor[mask < 255] = -1

So in python there actually is a map() function but usually there are better ways to do it (better in python; in other languages - like Haskell - map / fmap is obviously prefered in most contexts).

所以这里的关键点是首选方法是利用矢量化。 这也使代码更高效，因为这些张量操作是用低级语言实现的，而map()只不过是一个慢得多的 python-for 循环。