Pytorch boradcasting

I have a weight tensor(64,3, 3, 3) and a scale tensor(64), how to multiply a scale to per filter, a slow solution is showed below:

for i in range(weight.shape[0]):
    weight[i] = weight[i] * scale[i]

can boradcasting apply to this? please help me out, tks!

This should work:

q = torch.randn(64,3,3,3)
t = torch.randn(64)
qq = t.unsqueeze(1).unsqueeze(1).unsqueeze(1)  * q
for i in range(64):
    torch.allclose(t[i] * q[i], qq[i]) # should be true all the time

This should work:
weight = weight* scale.unsqueeze(dim=1).unsqueeze(dim=1).unsqueeze(dim=1)
or
weight = weight* scale.view(-1, 1, 1, 1)
etc.

1 Like

It works! thanks a lot.