argmax对softmax的输出返回概率最大的类别

最新推荐文章于 2023-09-30 21:35:58 发布

歌者And贰向箔

最新推荐文章于 2023-09-30 21:35:58 发布

阅读量4.7k

点赞数 3

分类专栏：工具文章标签：神经网络深度学习 pytorch 机器学习

本文链接：https://blog.csdn.net/ziqingnian/article/details/111272798

版权

工具专栏收录该内容

27 篇文章 2 订阅

订阅专栏

这篇博客介绍了如何处理softmax输出的概率分布，通过使用PyTorch的argmax函数来找出每行数据的最大概率对应的类别。首先，将numpy数组转换为torch tensor，然后指定dim=1来在行上求最大值索引。最后，将结果保存到文件中。

摘要由CSDN通过智能技术生成

假设数据如下，每行是softmax输出得到的概率，我需要找到最大的概率返回类别，可以使用argmax函数

(1)注意使用argmax函数时，需要将数据转换为tensor类型，否则报错
argmax(): argument 'input' (position 1) must be Tensor, not numpy.ndarray
(2)torch.argmax函数需要传递dim参数，dim=1就是在行上求
index = torch.argmax(data_pre, dim=1）

import numpy as np
import pandas as pd
import torch


data_pre = np.loadtxt('./pred.txt')
data_pre = torch.tensor(data_pre)
index = torch.argmax(data_pre, dim=1)
index = np.array(index)

np.savetxt('./cluster.txt', (index))