tensor.topk 以及tensor.argmax

tensor.topk 以及tensor.argmax

topk(self, k, dim=None, largest=True, sorted=True):,返回两个值,values与indices。

argmax(self, dim=None, keepdim=False): 返回Tensor

python 复制代码
'''
具体使用方法定义里写的很清楚,topk中largest=False返回最小值,sorted打乱原有的元素顺序。
argmax中keepdim=True:保持维数不变,默认会减少一维。
''' 

import torch

a = torch.randn(2, 2, 2)
a.argmax(1)

'''
tensor([[0, 1],
        [1, 0]])
'''

a.topk(k=1)
'''
torch.return_types.topk(
values=tensor([[[ 0.8766],
         [-0.1330]],

        [[ 1.5773],
         [ 0.8146]]]),
indices=tensor([[[0],
         [1]],

        [[1],
         [0]]]))
'''
a.topk(k=1).values
'''
tensor([[[ 0.8766],
         [-0.1330]],

        [[ 1.5773],
         [ 0.8146]]])
'''
a.topk(1).indices
'''
tensor([[[0],
         [1]],

        [[1],
         [0]]])
'''
相关推荐
强哥之神21 分钟前
英伟达发布 Llama Nemotron Nano 4B:专为边缘 AI 和科研任务优化的高效开源推理模型
人工智能·深度学习·语言模型·架构·llm·transformer·边缘计算
Green1Leaves23 分钟前
pytorch学习-9.多分类问题
人工智能·pytorch·学习
kyle~1 小时前
计算机视觉---RealSense深度相机技术
人工智能·数码相机·计算机视觉·机器人·嵌入式·ros·传感器
碣石潇湘无限路1 小时前
【AI篇】当Transformer模型开始学习《孙子兵法》
人工智能·学习
看到我,请让我去学习2 小时前
OpenCV开发-初始概念
人工智能·opencv·计算机视觉
汀沿河2 小时前
8.1 prefix Tunning与Prompt Tunning模型微调方法
linux·运维·服务器·人工智能
陈敬雷-充电了么-CEO兼CTO2 小时前
大模型技术原理 - 基于Transformer的预训练语言模型
人工智能·深度学习·语言模型·自然语言处理·chatgpt·aigc·transformer
学术 学术 Fun2 小时前
✨ OpenAudio S1:影视级文本转语音与语音克隆Mac整合包
人工智能·语音识别
风铃喵游3 小时前
让大模型调用MCP服务变得超级简单
前端·人工智能
旷世奇才李先生3 小时前
Pillow 安装使用教程
深度学习·microsoft·pillow