无脑入门pytorch系列(二)—— torch.mean

本系列教程适用于没有任何pytorch的同学(简单的python语法还是要的),从代码的表层出发挖掘代码的深层含义,理解具体的意思和内涵。pytorch的很多函数看着非常简单,但是其中包含了很多内容,不了解其中的意思就只能【看懂代码】,无法【理解代码】。

目录

官方定义

顾名思义,torch.mean返回输入张量中所有元素的平均值:

复制代码
def mean(input: Tensor, dim: Sequence[Union[str, ellipsis, None]], keepdim: _bool=False, *, dtype: Optional[_dtype]=None, out: Optional[Tensor]=None) -> Tensor:

上述的源码的定义,看起来参数十分多,但实际只需要记住下面几个参数:

  • input,输出是一个张量(tensor),注意如果不是tensor可以通过torch.tensor转换为tensor
  • dim,取平均值的维度,默认值是对tensor里的所有元素取平均值
  • keepdim,即保留张量的维度,因为取平均值后肯定是降维的,但是keepdim=True可以使得输出张量的维度与输入张量保持一致

官方的文档如下,torch.mean

demo

看下面一个例子:

python 复制代码
matrix = [[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]

tensor = torch.tensor(matrix)
print(tensor)

输出的结果:

直接求mean:

复制代码
torch.mean(tensor)

出现错误:RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long

意思是mean(),只能接受浮点数或复数类型的张量作为输入,所以先将tensor转化为float:

复制代码
tensor = tensor.float()

不配置任何参数

直接使用torch.mean():

复制代码
torch.mean(tensor)

输出的结果正好是1+2+...+9的平均值是5,所以如果不设置任何参数,那么默认取各个维度的平均值

设置按维度求平均

维度0:

复制代码
torch.mean(tensor, dim=0)

可以理解为矩阵按求平均值。

维度0:

复制代码
torch.mean(tensor, dim=0)

可以理解为矩阵按求平均值。

设置keepdim=True

复制代码
torch.mean(tensor, dim=0, keepdim=True)

其实就是在torch.mean(tensor, dim=0)的基础上,输出的一位张量上加上一对\[\]从而变为二维张量。

因为之前是按列求和,所以最后压缩为一行,然后补充行。

复制代码
torch.mean(tensor, dim=1, keepdim=True)

而dim=1是按行计算平均值,最后压缩的是列。

总结

torch.mean是个比较简单的函数,但是需要明白以下两点:

  • 第二个参数dim,决定了按哪个维度进行计算
  • 第三个参数keepdim,可以将输出张量的维度与输入张量保持一致
相关推荐
去码头整点薯条ing25 分钟前
某当网登录滑块【协议+OCR】
爬虫·python·ocr
zSD55rt5a32 分钟前
方差在扩散模型保护中的作用
人工智能·harmonyos
人生百态,人生如梦1 小时前
情感交互仿生人从技术到落地构想3——技术交流贴(2026.7)
人工智能·机器学习·人机交互·交互·具身智能
用户652238438111 小时前
为什么成熟的 LLM 应用,都把"后端地址"做成可切换的配置
人工智能
Aa99883341 小时前
AI视觉检测设备厂家的技术评估框架——以及AI视觉检测的检出边界与适用条件
人工智能
wuhanzhanhui1 小时前
当AI遇见储能!2026武汉国际数字能源产业与新型电力及储能技术展览会
人工智能·能源
赶紧写完去睡觉1 小时前
Anaconda 创建虚拟环境与使用
python·pycharm·conda
迷途呀1 小时前
Python:函数中的参数类型
开发语言·笔记·python·langchain·nlp
Tangyuewei1 小时前
给 AI 套上缰绳:Harness Engineering 是什么
人工智能
AI棒棒牛1 小时前
第11讲 | 目标检测任务:边界框、IoU、类别与置信度
人工智能·yolo·目标检测·yolo26