Векторно-тензорное поэлементное умножение в Pytorch

Question

Я пытаюсь извлечь яркость из тензора, представляющего изображение в Pytorch, и поэтому мне нужно умножить поэлементно вектор размера 3 (для трех весовых значений RGB) на тензор 3xNxN, представляющий изображение так, чтобы Я получаю матрицу NxN в конце, где три канала тензора суммируются с весами, указанными в векторе.

Я предполагаю, что существуют операции Pytorch, которые помогли бы мне сделать это без циклов, но у меня нет не нашел их.

Answer 1

Вы должны изменить свой 3 размерный вектор RGB, чтобы его можно было транслировать в 3xNxN следующим образом:

rgb = rgb.reshape(-1, 1, 1)

Таким образом, он будет иметь форму (3, 1, 1)

Теперь вы можете умножьте его на исходное изображение и суммируйте по первому измерению:

result = torch.sum(rgb * image, dim=0)