深入了解torch.arange()和torch.range()的用法及支持的数据类型格式

释放双眼,带上耳机,听听看~!
本文深入探讨了torch.arange()和torch.range()的用法及支持的数据类型格式,包括float类型在CPU、GPU、NPU上的支持情况,帮助读者更加全面地了解这两个函数。

本文已参与「新人创作礼」活动,一起开启掘金创作之路。

CPU(Central Processing Unit):中央处理器
GPU(Graphics Processing Unit):图形处理器
NPU(Neural Network Processing Unit):神经网络处理器,是基于神经网络算法与加速的新型处理器总称。

一、 torch.arange() 和 torch.range() 的用法

pytorch官网介绍:

深入了解torch.arange()和torch.range()的用法及支持的数据类型格式

torch.arange(start,end,step) 用于产生一个从start开始,到end结束(注意不包括end),步长为step的Tensor, 并且可以设置 Tensor 的 device 和 dtype

torch.arange 与 torch.range 功能及其相似,不同之处在于 torch.range(start,end,step) 生成的 Tensor, 包括 end

如:

a=torch.arange(1, 7, 2)
b=torch.range(1, 7, 2)
print(a)
print(b)

输出:

tensor([1, 3, 5])
tensor([1., 3., 5., 7.])

但是建议使用 torch.arange ,因为 torch.range 即将被pytorch 移除:

深入了解torch.arange()和torch.range()的用法及支持的数据类型格式

二、 torch.arange 支持的数据类型格式

只考虑 float 类型

cpu 不支持 float16,支持 float32 和 float64
cpu 支持 float16 、float32 和 float64
npu 不支持 float16 和 float64 ,只支持 float32

事实上 npu 基本不支持所有的 64位类型,包括 int64 和 float64,与算子无关。当然cpu 是支持 16位数据类型的,只是 torch.arange 不支持而已。

验证代码如下:

import torch

# CPU
a=torch.arange(1, 10, 2,device="cpu",dtype=torch.float16)    # 不可以
a=torch.arange(1, 10, 2,device="cpu",dtype=torch.float32)    # 可以
a=torch.arange(1, 10, 2,device="cpu",dtype=torch.float64)    # 可以

# GPU
a=torch.arange(1, 10, 2, device="cuda:0",dtype=torch.float16)    # 可以
a=torch.arange(1, 10, 2, device="cuda:0",dtype=torch.float32)    # 可以
a=torch.arange(1, 10, 2, device="cuda:0",dtype=torch.float64)    # 可以

# NPU
a=torch.arange(1, 10, 2, device="npu:0",dtype=torch.float16)    # 不可以
a=torch.arange(1, 10, 2, device="npu:0",dtype=torch.float32)    # 可以
a=torch.arange(1, 10, 2, device="npu:0",dtype=torch.float64)    # 不可以

print(a)
本网站的内容主要来自互联网上的各种资源,仅供参考和信息分享之用,不代表本网站拥有相关版权或知识产权。如您认为内容侵犯您的权益,请联系我们,我们将尽快采取行动,包括删除或更正。
AI教程

大型学术机构命名文本分类数据集LoT-insts介绍

2023-12-13 16:29:14

AI教程

PyTorch 中的 graph mode 和 Torch.FX 优化工具介绍

2023-12-13 16:37:14

个人中心
购物车
优惠劵
今日签到
有新私信 私信列表
搜索