无脑入门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,可以将输出张量的维度与输入张量保持一致
相关推荐
neocheng_52219 小时前
拒绝短视频浪费时间,自学AI实现自我提升
人工智能
LCG元19 小时前
深耕 RAG 工程,解锁大模型知识库开发之路
人工智能·语言模型
段一凡-华北理工大学19 小时前
2026 高炉炼铁智能化技术全景与演进路径~系列文章04:云-边-端协同架构:高炉智能化底层支撑体系
数据库·人工智能·深度学习·神经网络·安全·架构·高炉炼铁智能化
松☆19 小时前
ops-cv:昇腾NPU上的视觉算子,跟OpenCV有什么不一样?
人工智能·opencv·计算机视觉
2601_9578885620 小时前
短视频矩阵系统的AI调度引擎架构解析:从多平台API到智能分发的技术实现
人工智能·矩阵·音视频
m0_6872475220 小时前
AI绘画零基础:从文字到图像的完整指南
人工智能·ai作画·aigc
dxxt_yy20 小时前
通信对抗新利器:HWG1在铁路高速领域的卓越应用
人工智能·信息与通信
Rocky Ding*20 小时前
昔日AI绘画框架王者Stable Diffusion WebUI,已死
人工智能·深度学习·机器学习·chatgpt·ai作画·stable diffusion·aigc
localbob20 小时前
Anima二次元动漫必备大模型教程|20亿参数Anima模型+ComfyUI工作流下载 Anima模型下载、Anima动漫模型、AI绘画模型推荐、动漫AI绘画、ComfyUI工作流、AI游戏美术
人工智能·ai作画·游戏美术
2601_9577875820 小时前
2026年自媒体矩阵系统技术观察:当“人海战术“退场,AI如何重构内容分发逻辑?
人工智能·矩阵·媒体