如何对 PyTorch 张量进行哈希处理

Question

给定一个 PyTorch 整数张量，例如 torch.Size(N, C, H, W)。是否有一种有效的方法来哈希张量的每个元素，使得我可以获得从 [-MAX_INT32 to +MAX_INT32] 或 [0 to MAX_INT32] 的输出，快速在 GPU 上运行？

同时，我能够执行 output % N，并且每个元素将均匀或几乎均匀地分布在 0 - N 之间。

- Christian__

---

由于张量的每个元素都是一个整数，而整数本身是可哈希的（即可以用作字典键），那么为什么不能直接使用张量值作为哈希值呢？我觉得在这里需要更多关于你最终目标的上下文信息。 - undefined

如果你能提供一个你想要做的示例代码，我认为问题会更加清晰明了。 - undefined

为什么不只是用hash(tensor)呢？ - undefined

最简单的方法可能是使用线性同余生成器。它只是乘法、加法和取模运算。还有Thinc，它具有murmurhash3的cuda实现，但我还没有尝试过。 - undefined

@ValentinGoldité 这不满足属性 hash(tensor) == hash(tensor.clone())。也许 hash(tuple(tensor.reshape(-1).tolist())) 更好一些。 - undefined

Answer 1

不是GPU高效，也不是加密安全的，但是一个简单的方法是：

def hash_tensor(tensor):
    return hash(tuple(tensor.reshape(-1).tolist()))

测试：

>>> hash_tensor(torch.arange(8)) == hash_tensor(torch.arange(8))
True

>>> hash_tensor(torch.arange(8)) == hash_tensor(torch.arange(8) + 42)
False

- Mateen Ulhaq

---

请记住，除非您将PYTHONHASHSEED设置为已知值，否则hash输出在解释器重新运行之间是不稳定的。 - undefined

Answer 2

下面是一个基于线性同余生成器和musl/newlib参数的“粗略”并行化算法，适用于整数。

它通过反复减少最后一个维度，直到没有剩余的维度为止。因此，所需的在GPU上进行的“for循环”迭代总数为sum(x.shape)。如果你仔细选择形状（例如(8, 8, 8, ...)），这意味着只需要8 + 8 + 8 + ...次迭代。

import torch
from torch import Tensor


MULTIPLIER = 6364136223846793005
INCREMENT = 1
MODULUS = 2**64


def hash_tensor(x: Tensor) -> Tensor:
    assert x.dtype == torch.int64
    while x.ndim > 0:
        x = _reduce_last_axis(x)
    return x


@torch.no_grad()
def _reduce_last_axis(x: Tensor) -> Tensor:
    assert x.dtype == torch.int64
    acc = torch.zeros_like(x[..., 0])
    for i in range(x.shape[-1]):
        acc *= MULTIPLIER
        acc += INCREMENT
        acc += x[..., i]
        # acc %= MODULUS  # Not really necessary.
    return acc

测试：

x = torch.arange(8 * 8 * 8 * 8, dtype=torch.int64).reshape(8, 8, 8, 8)

>>> hash_tensor(x)
tensor(2150010819114838296)

>>> hash_tensor(x - 1)
tensor(2225417619311630616)

>>> hash_tensor(x * 2)
tensor(-4806624569712897768)

# Zeroes are OK:
>>> hash_tensor(torch.zeros((8, 8, 8, 8), dtype=torch.int64))
tensor(9106646207942574360)

# "Breaks" if you do a single axis of -1's...
>>> hash_tensor(torch.full((8,), -1))
tensor(0)

# But multiple axes is OKish:
>>> hash_tensor(torch.full((8, 8, 8, 8), -1))
tensor(9182053008139366680)

注意事项：

- 对于不可分割的张量长度，考虑在右侧填充零，直到它是（8, 8, 8, ...）的倍数（或者根据您想要并行化的程度）。 - 对于浮点数，可以考虑四舍五入（不准确），或者使用torch.frexp（准确）将其转换为整数。 - 由于字节顺序（即字节顺序）的不同，不同的设备/GPU可能会产生不同的结果。为了避免这种情况，您可能需要使用类似于[x >> 24, x >> 16, x >> 8, x >> 0]的方法将您的张量转换为正确顺序的字节流，然后进行扁平化处理。

- Mateen Ulhaq