Как мне сделать свертку в pytorch, где ядро ​​отличается для каждого мини-пакета?

Question

Позвольте мне express заголовок с примером:

Пусть A - тензор формы [16, 15 , 128, 128] (что означает [пакетный размер, каналы, высота, ширина])

Пусть B - тензор формы [16, 3 , 128, 128] (что означает [размер пакета, каналы, высота, ширина])

Я хочу вывести тензор формы [16, 5 , 128, 128] (что означает [размер пакета, каналы, высота, ширина])

Где i_th канал из 5 каналов вывода вычисляется умножением поэлементно B на срез i_th из 3 каналов A и их суммирование по измерению канала.

Как бы вы сделали эту операцию в pytorch?

Спасибо!

PD: Очень трудно express узнать, что я хочу от операции, если я не понял, пожалуйста, спросите меня, и я попытаюсь объяснить это

Answer 1

Я думаю, вы ищете torch.repeat_interleave, чтобы помочь вам "расширить" тензор B, чтобы иметь 15 каналов (5 групп из 3 входных каналов):

extB = torch.repeat_interleave(B, 3, dim=1)  # extend B to have 15 channels
m = A * extB  # element wise multiplication of A with the extended version of B
# using some reshaping and mean we can get the 5 output channels you want
out = m.view(m.shape[0], 5, 3, *m.shape[2:]).mean(dim=2)