快捷方式

张量属性

每个 torch.Tensor 都有一个 torch.dtype、一个 torch.device 和一个 torch.layout

torch.dtype

class torch.dtype

一个 torch.dtype 对象表示一个 torch.Tensor 的数据类型。PyTorch 有十二种不同的数据类型

数据类型

dtype

传统构造函数

32 位浮点数

torch.float32torch.float

torch.*.FloatTensor

64 位浮点数

torch.float64torch.double

torch.*.DoubleTensor

64 位复数

torch.complex64torch.cfloat

128 位复数

torch.complex128torch.cdouble

16 位浮点数 1

torch.float16torch.half

torch.*.HalfTensor

16 位浮点数 2

torch.bfloat16

torch.*.BFloat16Tensor

8 位无符号整数

torch.uint8

torch.*.ByteTensor

8 位有符号整数

torch.int8

torch.*.CharTensor

16 位有符号整数

torch.int16torch.short

torch.*.ShortTensor

32 位有符号整数

torch.int32torch.int

torch.*.IntTensor

64 位有符号整数

torch.int64torch.long

torch.*.LongTensor

布尔型

torch.bool

torch.*.BoolTensor

1

有时称为 binary16:使用 1 位符号位、5 位指数位和 10 位尾数位。在需要高精度时很有用。

2

有时称为 Brain Floating Point:使用 1 位符号位、8 位指数位和 7 位尾数位。在需要大范围时很有用,因为它与 float32 具有相同数量的指数位

要判断一个 torch.dtype 是否为浮点数据类型,可以使用 is_floating_point 属性,如果数据类型是浮点数据类型,该属性返回 True

要判断一个 torch.dtype 是否为复数数据类型,可以使用 is_complex 属性,如果数据类型是复数数据类型,该属性返回 True

当算术运算 (add, sub, div, mul) 的输入 dtype 不同时,我们会通过找到满足以下规则的最小 dtype 来进行类型提升:

  • 如果标量操作数的类型类别高于张量操作数(其中 复数 > 浮点数 > 整数 > 布尔值),则提升到足以容纳该类别所有标量操作数的大小类型。

  • 如果零维张量操作数的类别高于其他维度操作数,则提升到足以容纳该类别所有零维张量操作数的大小和类别类型。

  • 如果不存在更高类别的零维操作数,则提升到足以容纳所有维度操作数的大小和类别类型。

浮点标量操作数的 dtype 为 torch.get_default_dtype(),非布尔整数标量操作数的 dtype 为 torch.int64。与 numpy 不同,我们在确定操作数的最小 dtypes 时不会检查值。量化类型和复数类型尚不支持。

类型提升示例

>>> float_tensor = torch.ones(1, dtype=torch.float)
>>> double_tensor = torch.ones(1, dtype=torch.double)
>>> complex_float_tensor = torch.ones(1, dtype=torch.complex64)
>>> complex_double_tensor = torch.ones(1, dtype=torch.complex128)
>>> int_tensor = torch.ones(1, dtype=torch.int)
>>> long_tensor = torch.ones(1, dtype=torch.long)
>>> uint_tensor = torch.ones(1, dtype=torch.uint8)
>>> bool_tensor = torch.ones(1, dtype=torch.bool)
# zero-dim tensors
>>> long_zerodim = torch.tensor(1, dtype=torch.long)
>>> int_zerodim = torch.tensor(1, dtype=torch.int)

>>> torch.add(5, 5).dtype
torch.int64
# 5 is an int64, but does not have higher category than int_tensor so is not considered.
>>> (int_tensor + 5).dtype
torch.int32
>>> (int_tensor + long_zerodim).dtype
torch.int32
>>> (long_tensor + int_tensor).dtype
torch.int64
>>> (bool_tensor + long_tensor).dtype
torch.int64
>>> (bool_tensor + uint_tensor).dtype
torch.uint8
>>> (float_tensor + double_tensor).dtype
torch.float64
>>> (complex_float_tensor + complex_double_tensor).dtype
torch.complex128
>>> (bool_tensor + int_tensor).dtype
torch.int32
# Since long is a different kind than float, result dtype only needs to be large enough
# to hold the float.
>>> torch.add(long_tensor, float_tensor).dtype
torch.float32
指定算术运算的输出张量时,我们允许将其强制转换为其 dtype,但以下情况除外:
  • 整数输出张量不能接受浮点张量。

  • 布尔输出张量不能接受非布尔张量。

  • 非复数输出张量不能接受复数张量

强制转换示例

# allowed:
>>> float_tensor *= float_tensor
>>> float_tensor *= int_tensor
>>> float_tensor *= uint_tensor
>>> float_tensor *= bool_tensor
>>> float_tensor *= double_tensor
>>> int_tensor *= long_tensor
>>> int_tensor *= uint_tensor
>>> uint_tensor *= int_tensor

# disallowed (RuntimeError: result type can't be cast to the desired output type):
>>> int_tensor *= float_tensor
>>> bool_tensor *= int_tensor
>>> bool_tensor *= uint_tensor
>>> float_tensor *= complex_float_tensor

torch.device

class torch.device

一个 torch.device 对象表示分配 torch.Tensor 的设备或将要分配的设备。

`torch.device` 包含一个设备类型(最常见的是“cpu”或“cuda”,但也可能是“mps”“xpu”“xla”“meta”)以及该设备类型的可选设备序号。如果未指定设备序号,则此对象将始终表示该设备类型的当前设备,即使在调用 torch.cuda.set_device() 之后也是如此;例如,使用设备 'cuda' 构建的 torch.Tensor 等同于 'cuda:X',其中 X 是 torch.cuda.current_device() 的结果。

可以通过 Tensor.device 属性访问 torch.Tensor 的设备。

可以通过字符串或字符串加设备序号来构造 torch.device 对象

通过字符串

>>> torch.device('cuda:0')
device(type='cuda', index=0)

>>> torch.device('cpu')
device(type='cpu')

>>> torch.device('mps')
device(type='mps')

>>> torch.device('cuda')  # current cuda device
device(type='cuda')

通过字符串和设备序号

>>> torch.device('cuda', 0)
device(type='cuda', index=0)

>>> torch.device('mps', 0)
device(type='mps', index=0)

>>> torch.device('cpu', 0)
device(type='cpu', index=0)

设备对象也可以用作上下文管理器,以改变张量分配的默认设备

>>> with torch.device('cuda:1'):
...     r = torch.randn(2, 3)
>>> r.device
device(type='cuda', index=1)

如果工厂函数传递了明确的非 None 设备参数,此上下文管理器无效。要全局更改默认设备,另请参阅 torch.set_default_device()

警告

此函数会对每次 Python 调用 torch API(不仅仅是工厂函数)施加轻微的性能开销。如果这给您带来了问题,请在 https://github.com/pytorch/pytorch/issues/92701 上评论。

注意

函数中的 torch.device 参数通常可以用字符串代替。这使得代码快速原型设计成为可能。

>>> # Example of a function that takes in a torch.device
>>> cuda1 = torch.device('cuda:1')
>>> torch.randn((2,3), device=cuda1)
>>> # You can substitute the torch.device with a string
>>> torch.randn((2,3), device='cuda:1')

注意

出于历史原因,可以通过单个设备序号来构造设备,这被视为当前 加速器 类型。这与 Tensor.get_device() 匹配,后者返回设备张量的序号,但不支持 CPU 张量。

>>> torch.device(1)
device(type='cuda', index=1)

注意

接受设备参数的方法通常会接受(格式正确的)字符串或(传统的)整数设备序号,例如,以下所有写法都是等效的

>>> torch.randn((2,3), device=torch.device('cuda:1'))
>>> torch.randn((2,3), device='cuda:1')
>>> torch.randn((2,3), device=1)  # legacy

注意

张量永远不会在设备之间自动移动,需要用户显式调用。标量张量 (tensor.dim()==0) 是此规则的唯一例外,它们在需要时会自动从 CPU 传输到 GPU,因为此操作是“免费”的。示例

>>> # two scalars
>>> torch.ones(()) + torch.ones(()).cuda()  # OK, scalar auto-transferred from CPU to GPU
>>> torch.ones(()).cuda() + torch.ones(())  # OK, scalar auto-transferred from CPU to GPU
>>> # one scalar (CPU), one vector (GPU)
>>> torch.ones(()) + torch.ones(1).cuda()  # OK, scalar auto-transferred from CPU to GPU
>>> torch.ones(1).cuda() + torch.ones(())  # OK, scalar auto-transferred from CPU to GPU
>>> # one scalar (GPU), one vector (CPU)
>>> torch.ones(()).cuda() + torch.ones(1)  # Fail, scalar not auto-transferred from GPU to CPU and non-scalar not auto-transferred from CPU to GPU
>>> torch.ones(1) + torch.ones(()).cuda()  # Fail, scalar not auto-transferred from GPU to CPU and non-scalar not auto-transferred from CPU to GPU

torch.layout

class torch.layout

警告

`torch.layout` 类处于 beta 阶段,可能会发生变化。

一个 torch.layout 对象表示一个 torch.Tensor 的内存布局。目前,我们支持 torch.strided(密集张量),并对 torch.sparse_coo(稀疏 COO 张量)提供 beta 支持。

`torch.strided` 表示密集张量,是目前最常用的内存布局。每个 strided 张量都关联一个 torch.Storage,用于存储其数据。这些张量提供存储的多维、跨步视图。跨步(Strides)是一组整数:第 k 个跨步表示在张量的第 k 维中从一个元素移动到下一个元素所需的内存跳转。这个概念使得许多张量操作能够高效执行。

示例

>>> x = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
>>> x.stride()
(5, 1)

>>> x.t().stride()
(1, 5)

有关 torch.sparse_coo 张量的更多信息,请参阅 torch.sparse

torch.memory_format

class torch.memory_format

一个 torch.memory_format 对象表示分配 torch.Tensor 的内存格式或将要分配的内存格式。

可能的值包括

  • torch.contiguous_format:张量分配在或将分配在紧密不重叠的内存中。跨步值按递减顺序排列。

  • torch.channels_last:张量分配在或将分配在紧密不重叠的内存中。跨步值按 strides[0] > strides[2] > strides[3] > strides[1] == 1 的顺序排列,即 NHWC 顺序。

  • torch.channels_last_3d:张量分配在或将分配在紧密不重叠的内存中。跨步值按 strides[0] > strides[2] > strides[3] > strides[4] > strides[1] == 1 的顺序排列,即 NDHWC 顺序。

  • torch.preserve_format:用于像 clone 这样的函数中,以保留输入张量的内存格式。如果输入张量分配在紧密不重叠的内存中,则输出张量的跨步将从输入复制。否则,输出跨步将遵循 torch.contiguous_format

文档

访问 PyTorch 全面的开发者文档

查看文档

教程

获取针对初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源