python - 如何向量化总和？张量[i, :, :,:] + 张量[i]

Question

我想矢量化以下代码:

def style_noise(self, y, style):
    n = torch.randn(y.shape)
    for i in range(n.shape[0]):
        n[i] = (n[i] - n.mean(dim=(1, 2, 3))[i]) * style.std(dim=(1, 2, 3))[i] / n.std(dim=(1, 2, 3))[i] + style.mean(dim=(1, 2, 3))[i]
    noise = Variable(n, requires_grad=False).to(y.device)
    return noise

我没有找到这样做的好方法。

y 和 style 是 4d 张量，例如 style.shape = y.shape = [64, 3, 128, 128]。

我想返回噪声张量，noise.shape = [64, 3, 128, 128]。

如果问题不清楚，请在评论中告诉我。

Answer 1

您的用例正是 .mean 和 .std 方法带有 keepdim 参数的原因。您可以利用它来启用 broadcasting semantics为您矢量化事物:

def style_noise(self, y, style):
    n = torch.randn(y.shape)
    n_mean = n.mean(dim=(1, 2, 3), keepdim=True)
    n_std = n.std(dim=(1, 2, 3), keepdim=True)
    style_mean = style.mean(dim=(1, 2, 3), keepdim=True)
    style_std = style.std(dim=(1, 2, 3), keepdim=True)
    n = (n - n_mean) * style_std / n_std + style_mean
    noise = Variable(n, requires_grad=False).to(y.device)
    return noise