Изменение тензора Pytorch

Question

У меня есть тензор размера (24, 2, 224, 224) в Pytorch.

• 24 = размер партии
• 2 = матрицы, представляющие передний план и фон
• 224 = размер изображения по высоте
• 224 = ширина изображения размер

Это вывод CNN, который выполняет двоичную сегментацию. В каждой ячейке 2 матриц хранится вероятность того, что этот пиксель будет передним или задним: [n][0][h][w] + [n][1][h][w] = 1 для каждой координаты

Я хочу преобразовать его в тензор размера (24, 1, 224, 224). Значения в новом слое должны быть 0 или 1 в соответствии с матрицей, в которой вероятность была выше.

Как я могу это сделать? Какую функцию я должен использовать?

Answer 1

Использование torch.argmax() (для PyTorch +0,4):

prediction = torch.argmax(tensor, dim=1) # with 'dim' the considered dimension 
prediction = prediction.unsqueeze(1) # to reshape from (24, 224, 224) to (24, 1, 224, 224)

---

Если версия PyTorch ниже 0.4.0, можно использовать tensor.max(), которая возвращает как максимальные значения, так и их индексы (но которые не дифференцируются по значениям индекса):

_, prediction = tensor.max(dim=1)
prediction = prediction.unsqueeze(1) # to reshape from (24, 224, 224) to (24, 1, 224, 224)