python - 一个带条件的线性张量

Question

我有两个张量：

a = torch.nn.Parameter(torch.rand(7, requires_grad=True))
b = torch.randint(0,60, (20,))

是否有一个可以创建一个x大小为 20 （类似于“b”）的张量（称为它）有条件的一个衬里（或一种快速而简短的方法）？
IE

[b<4 use a[0], 4 <=b<12 use a[1], 12<=b<22 use a[2], <28, <38, <50, >50] for every b

因此，如果：

b = [12, 93, 54, 0...]

我希望我的新张量“x”是：

x = [a[2],a[6], a[6]...]

我将使用这个“x”张量来训练并需要反向传播和可学习的值

IE

loss = torch.rand(20) * x

loss.backward() ...

因此，如果其中一个a不在 x 中，我希望它不会改变。

标签： pythonpytorch

Answer 1

您可以对条件的乘法掩码求和：

x = a[0]*(b<4) + a[1]*((4<=b)&(b<12)) + a[2]*((12<=b)&(b<22)) + a[3]*((22<=b)&(b<28)) + a[4]*((28<=b)&(b<30)) + a[5]*((30<=b)&(b<50)) + a[6]*(b>=50)