Search code examples
pythonpytorchmatrix-multiplication

Matrix multiplication (element-wise) from numpy to Pytorch

I got two numpy arrays (image and and environment map), 

MatA
MatB

Both with shapes (256, 512, 3)

When I did the multiplication (element-wise) with numpy:

prod = np.multiply(MatA,MatB)

I got the wanted result (visualize via Pillow when turning back to Image)

But when I did it using pytorch, I got a really strange result(not even close to the aforementioned).

I did it with the following code:

MatATensor = transforms.ToTensor()(MatA)
MatBTensor = transforms.ToTensor()(MatB)

prodTensor = MatATensor * MatBTensor

For some reasons, the shape for both MatATensor and MatBtensor is 

torch.Size([3, 256, 512])

Same for the prodTensor too. When I tried to reshape to (256,512,3), I got an error.

Is there a way to get the same result?

I am new to pytorch, so any help would be appreciated.

Solution

  • If you read the documentation of transforms.ToTensor() you'll see this transformation does not only convert a numpy array to torch.FloatTensor, but also transpose its dimensions from HxWx3 to 3xHxW.
    To "undo" this you'll need to

     prodasNp = (prodTensor.permute(2, 0, 1) * 255).to(torch.uint8).numpy()

    See permute for more information.