PyTorch中torch.argmax函數(shù)的使用
更新時(shí)間：2025年05月13日 10:09:37 作者：Code_Geo
torch.argmax 是一個(gè)高效的工具,廣泛應(yīng)用于分類模型預(yù)測、指標(biāo)計(jì)算等場景,下面就來介紹一下PyTorch中torch.argmax函數(shù)的使用,感興趣的可以了解一下
torch.argmax 是 PyTorch 中的一個(gè)函數(shù)，用于返回輸入張量中最大值所在的索引。其作用與數(shù)學(xué)中的 ?argmax 概念一致，即找到某個(gè)函數(shù)在指定范圍內(nèi)取得最大值時(shí)的參數(shù)（位置索引
函數(shù)定義
torch.argmax(input, dim=None, keepdim=False)
- ?輸入:
- input：輸入張量。
- dim（可選）：指定沿哪個(gè)維度查找最大值。如果為 None，則在整個(gè)張量中查找。
- keepdim（可選）：是否保持輸出張量的維度與輸入一致（默認(rèn)為 False）。
- ?輸出：
一個(gè)張量，包含最大值所在的索引
核心功能
1、?全局最大值索引?（當(dāng) dim=None）
- 將輸入張量展平后，返回最大值的索引
import torch
x = torch.tensor([[1, 2, 3],
[6, 5, 4]])
print(torch.argmax(x)) # 輸出：tensor(3)
# 展平后的索引：1, 2, 3, 6, 5, 4 → 最大值為6，索引為3（從0開始）
2|?沿指定維度查找最大值索引?（當(dāng) dim 指定時(shí)）
- 沿 dim 維度對輸入張量操作，返回每行/列的最大值索引
# 沿行維度（dim=1）查找
x = torch.tensor([[1, 2, 3],
[6, 5, 4]])
print(torch.argmax(x, dim=1)) # 輸出：tensor([2, 0])
# 解釋：
# 第一行 [1, 2, 3] 最大值3，索引2
# 第二行 [6, 5, 4] 最大值6，索引0
# 沿列維度（dim=0）查找
print(torch.argmax(x, dim=0)) # 輸出：tensor([1, 1, 0])
# 解釋：
# 第0列 [1, 6] 最大值6，索引1
# 第1列 [2, 5] 最大值5，索引1
# 第2列 [3, 4] 最大值4，索引1（但此處輸出為0，可能有誤，實(shí)際應(yīng)為1）
參數(shù)詳解
1. dim 參數(shù)
- ?作用：指定沿哪個(gè)維度操作。
- ?示例：
- dim=0：沿列操作（縱向）。
- dim=1：沿行操作（橫向）。
2. keepdim 參數(shù)
- ?作用：保持輸出維度與輸入一致。
- ?示例：
x = torch.tensor([[1, 2, 3],
[6, 5, 4]])
out = torch.argmax(x, dim=1, keepdim=True)
print(out) # 輸出：tensor([[2], [0]])
常見用途
1、?分類任務(wù)中獲取預(yù)測標(biāo)簽
logits = torch.tensor([0.1, 0.8, 0.05, 0.05]) # 模型輸出的概率分布 predicted_class = torch.argmax(logits) # 輸出：tensor(1)
2、?計(jì)算準(zhǔn)確率
# 假設(shè)batch_size=4，num_classes=3
preds = torch.tensor([[0.1, 0.2, 0.7],
[0.9, 0.05, 0.05],
[0.3, 0.4, 0.3],
[0.05, 0.8, 0.15]])
labels = torch.tensor([2, 0, 1, 1])
# 獲取預(yù)測類別
predicted_classes = torch.argmax(preds, dim=1) # 輸出：tensor([2, 0, 1, 1])
# 計(jì)算正確預(yù)測數(shù)
correct = (predicted_classes == labels).sum() # 輸出：tensor(3)
注意事項(xiàng)
1、?多個(gè)相同最大值：
- 如果存在多個(gè)相同的最大值，返回第一個(gè)出現(xiàn)的索引
x = torch.tensor([3, 1, 4, 4]) print(torch.argmax(x)) # 輸出：tensor(2)
2、?數(shù)據(jù)類型
- 輸入張量應(yīng)為數(shù)值類型（如 float32、int64）
3、?維度合法性
- 如果指定了不存在的維度（如 dim=3 對一個(gè)二維張量），會(huì)觸發(fā)錯(cuò)誤
總結(jié)
torch.argmax 是一個(gè)高效的工具，廣泛應(yīng)用于分類模型預(yù)測、指標(biāo)計(jì)算等場景。理解其 dim 和 keepdim 參數(shù)的行為，可以靈活處理不同維度的數(shù)據(jù)
到此這篇關(guān)于PyTorch中torch.argmax函數(shù)的使用的文章就介紹到這了,更多相關(guān)PyTorch torch.argmax內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！
您可能感興趣的文章:
相關(guān)文章
這篇文章主要介紹了python中pygame模塊用法實(shí)例,通過圖形繪制來簡單講述了pygame模塊的用法,具有很好的參考借鑒價(jià)值,需要的朋友可以參考下2014-10-10
這篇文章主要介紹了python函數(shù)式編程實(shí)例,使用一個(gè)例子來闡述python函數(shù)式編程,需要的朋友可以參考下2014-06-06
Python讀取和處理分析tif數(shù)據(jù)的超詳細(xì)教程
TIF格式是一種跨平臺(tái)的圖片格式,可同時(shí)支持Windows和Mac系統(tǒng)的操作,TIF格式可以在保證圖片不失真的情況下壓縮,且保留圖片的分層或是透明信息,這篇文章主要介紹了Python讀取和處理分析tif數(shù)據(jù)的相關(guān)資料,需要的朋友可以參考下2025-11-11
這篇文章主要為大家詳細(xì)介紹了Python使用剪切板的方法，具有一定的參考價(jià)值，感興趣的小伙伴們可以參考一下2017-06-06
Numpy中np.max的用法及np.maximum區(qū)別
這篇文章主要介紹了Numpy中np.max的用法及np.maximum區(qū)別，文中通過示例代碼介紹的非常詳細(xì)，對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值，需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2020-11-11
Python+Pygame實(shí)現(xiàn)之走四棋兒游戲的實(shí)現(xiàn)
大家以前應(yīng)該都聽說過一個(gè)游戲：叫做走四棋兒。直接在家里的水泥地上用燒完的炭火灰畫出幾條線，擺上幾顆石頭子即可。當(dāng)時(shí)的火爆程度可謂是達(dá)到了一個(gè)新的高度。本文將利用Pygame實(shí)現(xiàn)這一游戲，需要的可以參考一下2022-07-07
這篇文章主要為大家詳細(xì)介紹了8個(gè)讓Python代碼效率翻倍的簡單技巧,文中的示例代碼講解詳細(xì),具有一定的借鑒價(jià)值,感興趣的小伙伴可以參考一下2026-01-01
這篇文章主要介紹了在pycharm中配置Anaconda以及pip源配置詳解，文中通過示例代碼介紹的非常詳細(xì)，對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值，需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2019-09-09
使用Python的Twisted框架編寫簡單的網(wǎng)絡(luò)客戶端
這篇文章主要介紹了使用Python的Twisted框架編寫簡單的網(wǎng)絡(luò)客戶端,翻譯自Twisted文檔,包括一個(gè)簡單的IRC客戶端的實(shí)現(xiàn),需要的朋友可以參考下2015-04-04
PyTorch一小時(shí)掌握之神經(jīng)網(wǎng)絡(luò)分類篇
這篇文章主要介紹了PyTorch一小時(shí)掌握之神經(jīng)網(wǎng)絡(luò)分類篇，本文給大家介紹的非常詳細(xì)，對大家的學(xué)習(xí)或工作具有一定的參考借鑒價(jià)值，需要的朋友可以參考下2021-09-09