Пакетно-матричное умножение в Pytorch - путается с обработкой выходного измерения

Question

У меня есть два массива:

A
B

Массив A содержит серию изображений RGB с формой:

[batch, Width, Height, 3]

, тогда как Array B содержит коэффициенты, необходимые для "подобной преобразованию" операции с изображениями, с формой:

[batch, 4, 4, 3]

Проще говоря, операция для одного изображения - это умножение, которое выводит карту окружения (normalMap * Coefficients).

Вывод, который я хочу, должен содержать форму:

[batch, Width, Height, 3]

Я пытался использовать torch.bmm, но не получилось. Возможно ли это как-то?

Answer 1

Я думаю, вам нужно рассчитать, что PyTorch работает с форматом

BxCxHxW : number of mini-batches, channels, height, width

, а также использовать matmul, поскольку bmm работает с 3Dтензор.

Я знаю, вы можете найти это онлайн, но для любого случая:

batch1 = torch.randn(10, 3, 20, 10)
batch2 = torch.randn(10, 3, 10, 30)
res = torch.matmul(batch1, batch2)
res.size() # torch.Size([10, 3, 20, 30])