Torch.argmax(input, dim=None, keepdim=False)
argmax-Funktion: Gibt den Index des Maximalwerts der angegebenen Dimension zurück. Wenn dim eine bestimmte Dimension angibt, verschwindet diese Dimension und alle zurückgegebenen Dimensionen sind kleiner als die durch dim angegebene Dimension. Bestimmen Sie anhand der zurückgegebenen Dimension, welche Dimension, um den Argmax-Vorgang auszuführen
Wenn die Eingabe beispielsweise token_output ist, lauten die Dimensionen (62, 320, 523): target_len: 62 [maximale Sequenzlänge], 320 [Stapelgröße], 523 [Vokabulargröße]
output_all_token_id = torch.argmax(token_output, -1).tolist()Die Bedeutung dieser Passage besteht darin, die letzte Dimension verschwinden zu lassen (nehmen Sie das Token mit der höchsten Generierungswahrscheinlichkeit in jedem Stapel), dann wird daraus eine (62.320) Dimension, was 320 generierte Texte bedeutet
Einfaches Beispiel:
Wenn es sich um eine zweidimensionale Matrix handelt:
dim = 0 bedeutet, dass die Dimension von „Zeile“ verschwindet und nur Spalten übrig bleiben, d. h. der Index des Maximalwerts in jeder Spalte gesucht wird
dim = 1 bedeutet, dass die Dimension der „Spalte“ verschwindet und nur Zeilen übrig bleiben, d. h. der Index des Maximalwerts in jeder Zeile ermittelt wird
import torch
a = torch.randn(2, 3)
print(a)Tensor([[-0.3018, 0.3350, 0.8318],
[ 0.2485, 0.5349, -1.2342]])
# 求所有值中最大值的索引
print(torch.argmax(a))
# dim=0意思就是“行”这一维度消失，只剩下列，也就是求每一列中最大值的索引
print(torch.argmax(a, dim=0))
# dim=1意思就是“列”这一维度消失，只剩下行，也就是求每一行中最大值的索引
print(torch.argmax(a, dim=1))Tensor(2)
Tensor([2, 1])
Tensor([1, 1, 0])
Beschreibung der Funktion Torch.argmax_Egozjuers Blog-CSDN-Blog