Pytorch 批量逐行应用函数答案

【问题标题】：Pytorch batch row-wise application of functionPytorch 批量逐行应用函数
【发布时间】：2021-09-10 04:55:45
【问题描述】：

我想找到一种方法来应用一个计算成对距离的函数，我们称之为 dists(A, B)，对批处理中的每个输入元素按行表示，意思是：

(100, 16, 3) -- 输入，100 是批量大小，因此 100 个实例，16 是图像大小，3 个过滤器（要求 Conv2D）

(5, 3) -- 我要计算行距离的张量（假设它是 dists(A, B) 中的 A 并且是固定的）

现在，对于每个实例，我都应该返回一个形状为 (5, 16) 的矩阵。自然，我可以使用 for 来跨越批次并获得最终的 (100,5,16) 结果。不过，我很想知道是否有更简单的方法可以使用 GPU 以行方式并行应用我的函数。

非常感谢您的宝贵时间。

【问题讨论】：

标签： python machine-learning deep-learning parallel-processing pytorch

【解决方案1】：

假设我们使用的是 L1 距离：

import torch
# data and target
a = torch.randn(100, 16, 3)
b = torch.randn(5, 3)

# Reshape the tensors
a = a.unsqueeze(1)
b = b.unsqueeze(0).unsqueeze(2)

print(a.shape, b.shape)

# Compute distance
dist = (a-b).abs().sum(3)
print(dist.shape)

【讨论】：

好吧，我们不应该在 cmets 中回复谢谢，但是，嘿，非常感谢。现在我可以将整个过程加快很多。