как получить индексы строк и столбцов максимального элемента из тензора 2D-питора?

Question

Есть ли способ, которым я могу получить индексы строк и столбцов наибольшего элемента, содержащегося в 2-мерном тензоре Pytorch? Например, см. Тензор Pytorch a ниже:

a
>> torch.tensor([1,2,3],
                [9,5,4],
                [6,7,8])

Наибольший элемент в тензоре a равен 9, что происходит в первом столбце второго ряда. Если я изменю это на python индекс столбца и строки, который начинается с нуля, индекс столбца элемента будет 0, а индекс строки будет 1.

Есть ли способ получить индекс [1,0] из 2-мерного тензора Pytorch a?

Спасибо,

Answer 1

К сожалению, нет встроенного метода. Однако вы можете использовать numpy:

np.unravel_index(torch.argmax(a), a.shape)

В противном случае вам нужно написать собственный лог c, например:

def unravel_index(flat_idx, shape): 
     multi_idx = [] 
     r = flat_idx 
     for s in shape[:-1]: 
         multi_idx.append(r // s) 
         r = r % s 
     multi_idx.append(r % s) 
     return multi_idx