评价此页

介绍 || 张量 || Autograd || 构建模型 || TensorBoard 支持 || 训练模型 || 模型理解

PyTorch 张量入门#

创建日期:2021 年 11 月 30 日 | 最后更新:2025 年 9 月 22 日 | 最后验证:2024 年 11 月 05 日

请观看下面的视频,或在 youtube 上观看。

张量是 PyTorch 中的核心数据抽象。本交互式 Notebook 提供了对 torch.Tensor 类的深入介绍。

首先,让我们导入 PyTorch 模块。我们还将添加 Python 的 math 模块以方便某些示例。

import torch
import math

创建张量#

创建张量的最简单方法是调用 torch.empty()

x = torch.empty(3, 4)
print(type(x))
print(x)
<class 'torch.Tensor'>
tensor([[ 1.9192e-04,  4.5804e-41,  2.2958e-32,  3.0645e-41],
        [-4.3165e+09,  3.0644e-41, -1.9917e+27,  4.5801e-41],
        [-1.9896e+27,  4.5801e-41, -1.9917e+27,  4.5801e-41]])

让我们解析一下我们刚才所做的。

  • 我们使用 torch 模块附带的众多工厂方法之一创建了一个张量。

  • 张量本身是 2 维的,有 3 行 4 列。

  • 返回对象的类型是 torch.Tensor,它是 torch.FloatTensor 的别名;默认情况下,PyTorch 张量填充 32 位浮点数。(关于数据类型的更多信息见下文。)

  • 打印张量时,您可能会看到一些看起来随机的值。 torch.empty() 调用为张量分配内存,但并未用任何值对其进行初始化 - 因此您看到的是分配时内存中存在的任何内容。

关于张量及其维数和术语的简要说明。

  • 您有时会看到 1 维张量被称为向量

  • 同样,2 维张量通常被称为矩阵

  • 任何超过两维的通常都只称为张量。

十有八九,您会想用某个值来初始化张量。常见的情况是全零、全一或随机值,并且 torch 模块为所有这些提供了工厂方法。

zeros = torch.zeros(2, 3)
print(zeros)

ones = torch.ones(2, 3)
print(ones)

torch.manual_seed(1729)
random = torch.rand(2, 3)
print(random)
tensor([[0., 0., 0.],
        [0., 0., 0.]])
tensor([[1., 1., 1.],
        [1., 1., 1.]])
tensor([[0.3126, 0.3791, 0.3087],
        [0.0736, 0.4216, 0.0691]])

工厂方法都如您所料 - 我们有一个全零的张量,一个全一的张量,以及一个包含 0 到 1 之间随机值的张量。

随机张量和种子#

说到随机张量,您注意到它前面紧跟着 torch.manual_seed() 的调用了吗?用随机值初始化张量(例如模型的学习权重)很常见,但在某些情况下(尤其是在研究环境中),您会希望确保结果的可复现性。手动设置随机数生成器的种子是做到这一点的方法。让我们仔细看看。

tensor([[0.3126, 0.3791, 0.3087],
        [0.0736, 0.4216, 0.0691]])
tensor([[0.2332, 0.4047, 0.2162],
        [0.9927, 0.4128, 0.5938]])
tensor([[0.3126, 0.3791, 0.3087],
        [0.0736, 0.4216, 0.0691]])
tensor([[0.2332, 0.4047, 0.2162],
        [0.9927, 0.4128, 0.5938]])

您应该在上面看到 random1random3 具有相同的值,random2random4 也一样。手动设置 RNG 的种子会重置它,因此在大多数情况下,依赖随机数的相同计算应提供相同的结果。

有关更多信息,请参阅 PyTorch 关于可复现性的文档

张量形状#

通常,在对两个或多个张量执行操作时,它们需要具有相同的形状 - 即,具有相同的维数并在每个维数中具有相同的单元数。为此,我们有 torch.*_like() 方法。

torch.Size([2, 2, 3])
tensor([[[1.8992e-34, 3.0645e-41, 6.3185e-16],
         [4.5803e-41, 1.4013e-45, 0.0000e+00]],

        [[8.4078e-45, 0.0000e+00, 0.0000e+00],
         [0.0000e+00, 0.0000e+00, 0.0000e+00]]])
torch.Size([2, 2, 3])
tensor([[[2.2792e-32, 3.0645e-41, 9.1477e-41],
         [0.0000e+00, 1.6816e-44, 0.0000e+00]],

        [[3.5873e-43, 0.0000e+00, 6.6220e-06],
         [4.5804e-41, 0.0000e+00, 0.0000e+00]]])
torch.Size([2, 2, 3])
tensor([[[0., 0., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [0., 0., 0.]]])
torch.Size([2, 2, 3])
tensor([[[1., 1., 1.],
         [1., 1., 1.]],

        [[1., 1., 1.],
         [1., 1., 1.]]])
torch.Size([2, 2, 3])
tensor([[[0.6128, 0.1519, 0.0453],
         [0.5035, 0.9978, 0.3884]],

        [[0.6929, 0.1703, 0.1384],
         [0.4759, 0.7481, 0.0361]]])

代码单元格中的第一个新内容是使用张量上的 .shape 属性。此属性包含张量每个维度的范围列表 - 在我们的例子中,x 是一个 3 维张量,形状为 2 x 2 x 3。

在此下方,我们调用了 .empty_like().zeros_like().ones_like().rand_like() 方法。使用 .shape 属性,我们可以验证这些方法中的每一个都返回具有相同维度和范围的张量。

我们将介绍的最后一种创建张量的方法是从 PyTorch 集合直接指定其数据。

some_constants = torch.tensor([[3.1415926, 2.71828], [1.61803, 0.0072897]])
print(some_constants)

some_integers = torch.tensor((2, 3, 5, 7, 11, 13, 17, 19))
print(some_integers)

more_integers = torch.tensor(((2, 4, 6), [3, 6, 9]))
print(more_integers)
tensor([[3.1416, 2.7183],
        [1.6180, 0.0073]])
tensor([ 2,  3,  5,  7, 11, 13, 17, 19])
tensor([[2, 4, 6],
        [3, 6, 9]])

使用 torch.tensor() 是创建张量的最直接方法,如果您已经拥有 Python 元组或列表中的数据。如上所示,嵌套集合将导致多维张量。

注意

torch.tensor() 会创建数据的副本。

张量数据类型#

设置张量的数据类型可以通过几种方式实现。

a = torch.ones((2, 3), dtype=torch.int16)
print(a)

b = torch.rand((2, 3), dtype=torch.float64) * 20.
print(b)

c = b.to(torch.int32)
print(c)
tensor([[1, 1, 1],
        [1, 1, 1]], dtype=torch.int16)
tensor([[ 0.9956,  1.4148,  5.8364],
        [11.2406, 11.2083, 11.6692]], dtype=torch.float64)
tensor([[ 0,  1,  5],
        [11, 11, 11]], dtype=torch.int32)

设置张量的底层数据类型的最简单方法是在创建时使用一个可选参数。在上面的单元格的第一行,我们为张量 a 设置了 dtype=torch.int16。当我们打印 a 时,我们可以看到它填充的是 1 而不是 1. - 这是 Python 中表示整数类型而不是浮点数的细微提示。

打印 a 时还应注意的另一件事是,与我们保留 dtype 为默认值(32 位浮点数)时不同,打印张量时还会指定其 dtype

您可能还注意到,我们从指定张量的形状作为一系列整数参数,转变为将这些参数分组到元组中。这并非严格必要 - PyTorch 会将一系列初始的、未标记的整数参数视为张量形状 - 但在添加可选参数时,可以使您的意图更具可读性。

设置数据类型的另一种方法是使用 .to() 方法。在上面的单元格中,我们以通常方式创建了一个随机浮点张量 b。之后,我们使用 .to() 方法将 b 转换为 32 位整数,从而创建了 c。请注意,c 包含与 b 相同的值,但已截断为整数。

有关更多信息,请参阅 数据类型文档

PyTorch 张量的数学与逻辑运算#

既然您已经了解了一些创建张量的方法... 您可以用它们做什么?

让我们先看看基本算术,以及张量如何与简单的标量交互。

ones = torch.zeros(2, 2) + 1
twos = torch.ones(2, 2) * 2
threes = (torch.ones(2, 2) * 7 - 1) / 2
fours = twos ** 2
sqrt2s = twos ** 0.5

print(ones)
print(twos)
print(threes)
print(fours)
print(sqrt2s)
tensor([[1., 1.],
        [1., 1.]])
tensor([[2., 2.],
        [2., 2.]])
tensor([[3., 3.],
        [3., 3.]])
tensor([[4., 4.],
        [4., 4.]])
tensor([[1.4142, 1.4142],
        [1.4142, 1.4142]])

如上所示,张量与标量之间的算术运算(如加法、减法、乘法、除法和指数运算)会应用于张量的每个元素。因为此类运算的输出将是一个张量,所以您可以按照通常的运算符优先级规则将它们链接在一起,就像我们在创建 threes 的行中所做的那样。

两个张量之间的类似运算也如您直观预期那样工作。

powers2 = twos ** torch.tensor([[1, 2], [3, 4]])
print(powers2)

fives = ones + fours
print(fives)

dozens = threes * fours
print(dozens)
tensor([[ 2.,  4.],
        [ 8., 16.]])
tensor([[5., 5.],
        [5., 5.]])
tensor([[12., 12.],
        [12., 12.]])

这里需要注意的是,前一个代码单元格中的所有张量都具有相同的形状。当尝试对不同形状的张量执行二元运算时会发生什么?

注意

下面的单元格会引发运行时错误。这是故意的。

a = torch.rand(2, 3)
b = torch.rand(3, 2)

print(a * b)

在一般情况下,您不能以这种方式操作不同形状的张量,即使在像上面的单元格那样张量具有相同数量元素的情况下也不行。

简而言之:张量广播#

注意

如果您熟悉 NumPy ndarray 中的广播语义,您会发现这里也适用相同的规则。

相同形状规则的例外是张量广播。以下是一个示例。

rand = torch.rand(2, 4)
doubled = rand * (torch.ones(1, 4) * 2)

print(rand)
print(doubled)
tensor([[0.6146, 0.5999, 0.5013, 0.9397],
        [0.8656, 0.5207, 0.6865, 0.3614]])
tensor([[1.2291, 1.1998, 1.0026, 1.8793],
        [1.7312, 1.0413, 1.3730, 0.7228]])

这里的诀窍是什么?我们如何将一个 2x4 的张量乘以一个 1x4 的张量?

广播是一种对形状相似的张量执行运算的方法。在上面的示例中,一行四列的张量被乘以两行四列张量的两行

这是深度学习中的一项重要运算。常见示例是将学习权重的张量乘以输入张量的批次,分别对批次中的每个实例应用该运算,并返回一个形状相同的张量 - 就像我们上面的 (2, 4) * (1, 4) 示例返回一个形状为 (2, 4) 的张量一样。

广播规则是:

  • 每个张量至少必须有一个维度 - 没有空张量。

  • 比较两个张量的维度大小,从最后一个到第一个:

    • 每个维度必须相等,或者

    • 其中一个维度的大小必须是 1,或者

    • 其中一个张量中不存在该维度。

形状相同的张量当然是“可广播的”,正如您之前所见。

以下是一些符合上述规则并允许广播的情况示例:

a =     torch.ones(4, 3, 2)

b = a * torch.rand(   3, 2) # 3rd & 2nd dims identical to a, dim 1 absent
print(b)

c = a * torch.rand(   3, 1) # 3rd dim = 1, 2nd dim identical to a
print(c)

d = a * torch.rand(   1, 2) # 3rd dim identical to a, 2nd dim = 1
print(d)
tensor([[[0.6493, 0.2633],
         [0.4762, 0.0548],
         [0.2024, 0.5731]],

        [[0.6493, 0.2633],
         [0.4762, 0.0548],
         [0.2024, 0.5731]],

        [[0.6493, 0.2633],
         [0.4762, 0.0548],
         [0.2024, 0.5731]],

        [[0.6493, 0.2633],
         [0.4762, 0.0548],
         [0.2024, 0.5731]]])
tensor([[[0.7191, 0.7191],
         [0.4067, 0.4067],
         [0.7301, 0.7301]],

        [[0.7191, 0.7191],
         [0.4067, 0.4067],
         [0.7301, 0.7301]],

        [[0.7191, 0.7191],
         [0.4067, 0.4067],
         [0.7301, 0.7301]],

        [[0.7191, 0.7191],
         [0.4067, 0.4067],
         [0.7301, 0.7301]]])
tensor([[[0.6276, 0.7357],
         [0.6276, 0.7357],
         [0.6276, 0.7357]],

        [[0.6276, 0.7357],
         [0.6276, 0.7357],
         [0.6276, 0.7357]],

        [[0.6276, 0.7357],
         [0.6276, 0.7357],
         [0.6276, 0.7357]],

        [[0.6276, 0.7357],
         [0.6276, 0.7357],
         [0.6276, 0.7357]]])

仔细查看上面每个张量的值。

  • 创建 b 的乘法运算被广播到 a 的每个“层”。

  • 对于 c,该运算被广播到 a 的每个层和行 - 每个 3 元素列都相同。

  • 对于 d,我们颠倒了过来 - 现在每都相同,跨越层和列。

有关广播的更多信息,请参阅 PyTorch 文档

以下是一些尝试广播但会失败的示例:

注意

下面的单元格会引发运行时错误。这是故意的。

a =     torch.ones(4, 3, 2)

b = a * torch.rand(4, 3)    # dimensions must match last-to-first

c = a * torch.rand(   2, 3) # both 3rd & 2nd dims different

d = a * torch.rand((0, ))   # can't broadcast with an empty tensor

张量的更多数学运算#

PyTorch 张量有三百多种可对其执行的操作。

以下是来自主要操作类别的一些小样本。

# common functions
a = torch.rand(2, 4) * 2 - 1
print('Common functions:')
print(torch.abs(a))
print(torch.ceil(a))
print(torch.floor(a))
print(torch.clamp(a, -0.5, 0.5))

# trigonometric functions and their inverses
angles = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
sines = torch.sin(angles)
inverses = torch.asin(sines)
print('\nSine and arcsine:')
print(angles)
print(sines)
print(inverses)

# bitwise operations
print('\nBitwise XOR:')
b = torch.tensor([1, 5, 11])
c = torch.tensor([2, 7, 10])
print(torch.bitwise_xor(b, c))

# comparisons:
print('\nBroadcasted, element-wise equality comparison:')
d = torch.tensor([[1., 2.], [3., 4.]])
e = torch.ones(1, 2)  # many comparison ops support broadcasting!
print(torch.eq(d, e)) # returns a tensor of type bool

# reductions:
print('\nReduction ops:')
print(torch.max(d))        # returns a single-element tensor
print(torch.max(d).item()) # extracts the value from the returned tensor
print(torch.mean(d))       # average
print(torch.std(d))        # standard deviation
print(torch.prod(d))       # product of all numbers
print(torch.unique(torch.tensor([1, 2, 1, 2, 1, 2]))) # filter unique elements

# vector and linear algebra operations
v1 = torch.tensor([1., 0., 0.])         # x unit vector
v2 = torch.tensor([0., 1., 0.])         # y unit vector
m1 = torch.rand(2, 2)                   # random matrix
m2 = torch.tensor([[3., 0.], [0., 3.]]) # three times identity matrix

print('\nVectors & Matrices:')
print(torch.linalg.cross(v2, v1)) # negative of z unit vector (v1 x v2 == -v2 x v1)
print(m1)
m3 = torch.linalg.matmul(m1, m2)
print(m3)                  # 3 times m1
print(torch.linalg.svd(m3))       # singular value decomposition
Common functions:
tensor([[0.9238, 0.5724, 0.0791, 0.2629],
        [0.1986, 0.4439, 0.6434, 0.4776]])
tensor([[-0., -0., 1., -0.],
        [-0., 1., 1., -0.]])
tensor([[-1., -1.,  0., -1.],
        [-1.,  0.,  0., -1.]])
tensor([[-0.5000, -0.5000,  0.0791, -0.2629],
        [-0.1986,  0.4439,  0.5000, -0.4776]])

Sine and arcsine:
tensor([0.0000, 0.7854, 1.5708, 2.3562])
tensor([0.0000, 0.7071, 1.0000, 0.7071])
tensor([0.0000, 0.7854, 1.5708, 0.7854])

Bitwise XOR:
tensor([3, 2, 1])

Broadcasted, element-wise equality comparison:
tensor([[ True, False],
        [False, False]])

Reduction ops:
tensor(4.)
4.0
tensor(2.5000)
tensor(1.2910)
tensor(24.)
tensor([1, 2])

Vectors & Matrices:
tensor([ 0.,  0., -1.])
tensor([[0.7375, 0.8328],
        [0.8444, 0.2941]])
tensor([[2.2125, 2.4985],
        [2.5332, 0.8822]])
torch.return_types.linalg_svd(
U=tensor([[-0.7889, -0.6145],
        [-0.6145,  0.7889]]),
S=tensor([4.1498, 1.0548]),
Vh=tensor([[-0.7957, -0.6056],
        [ 0.6056, -0.7957]]))

这是一个小样本操作。有关更多详细信息和完整的数学函数列表,请查看 文档。有关更多详细信息和完整的线性代数运算列表,请查看此 文档

原地修改张量#

大多数张量的二元运算会返回第三个新张量。当我们说 c = a * b(其中 ab 是张量)时,新张量 c 将占据与其他张量不同的内存区域。

但是,有时您可能希望原地修改张量 - 例如,如果您正在进行一项元素级计算,可以丢弃中间值。为此,大多数数学函数都有一个带有附加下划线(_)的版本,该版本会原地修改张量。

例如

a = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('a:')
print(a)
print(torch.sin(a))   # this operation creates a new tensor in memory
print(a)              # a has not changed

b = torch.tensor([0, math.pi / 4, math.pi / 2, 3 * math.pi / 4])
print('\nb:')
print(b)
print(torch.sin_(b))  # note the underscore
print(b)              # b has changed
a:
tensor([0.0000, 0.7854, 1.5708, 2.3562])
tensor([0.0000, 0.7071, 1.0000, 0.7071])
tensor([0.0000, 0.7854, 1.5708, 2.3562])

b:
tensor([0.0000, 0.7854, 1.5708, 2.3562])
tensor([0.0000, 0.7071, 1.0000, 0.7071])
tensor([0.0000, 0.7071, 1.0000, 0.7071])

对于算术运算,有一些功能行为类似。

a = torch.ones(2, 2)
b = torch.rand(2, 2)

print('Before:')
print(a)
print(b)
print('\nAfter adding:')
print(a.add_(b))
print(a)
print(b)
print('\nAfter multiplying')
print(b.mul_(b))
print(b)
Before:
tensor([[1., 1.],
        [1., 1.]])
tensor([[0.3788, 0.4567],
        [0.0649, 0.6677]])

After adding:
tensor([[1.3788, 1.4567],
        [1.0649, 1.6677]])
tensor([[1.3788, 1.4567],
        [1.0649, 1.6677]])
tensor([[0.3788, 0.4567],
        [0.0649, 0.6677]])

After multiplying
tensor([[0.1435, 0.2086],
        [0.0042, 0.4459]])
tensor([[0.1435, 0.2086],
        [0.0042, 0.4459]])

请注意,这些原地算术函数是 torch.Tensor 对象上的方法,而不是像许多其他函数(例如 torch.sin())那样附加到 torch 模块。从 a.add_(b) 可以看出,调用张量是被原地更改的那个。

还有另一种选择可以将计算结果放置在现有的、已分配的张量中。我们到目前为止看到的许多方法和函数 - 包括创建方法!- 都有一个 out 参数,允许您指定一个张量来接收输出。如果 out 张量的形状和 dtype 正确,则可以在不进行新内存分配的情况下完成。

a = torch.rand(2, 2)
b = torch.rand(2, 2)
c = torch.zeros(2, 2)
old_id = id(c)

print(c)
d = torch.matmul(a, b, out=c)
print(c)                # contents of c have changed

assert c is d           # test c & d are same object, not just containing equal values
assert id(c) == old_id  # make sure that our new c is the same object as the old one

torch.rand(2, 2, out=c) # works for creation too!
print(c)                # c has changed again
assert id(c) == old_id  # still the same object!
tensor([[0., 0.],
        [0., 0.]])
tensor([[0.3653, 0.8699],
        [0.2364, 0.3604]])
tensor([[0.0776, 0.4004],
        [0.9877, 0.0352]])

复制张量#

与 Python 中的任何对象一样,将张量分配给变量会使该变量成为张量的标签,而不会复制它。例如:

a = torch.ones(2, 2)
b = a

a[0][1] = 561  # we change a...
print(b)       # ...and b is also altered
tensor([[  1., 561.],
        [  1.,   1.]])

但是,如果您想有一个单独的数据副本进行操作怎么办?clone() 方法可供您使用。

a = torch.ones(2, 2)
b = a.clone()

assert b is not a      # different objects in memory...
print(torch.eq(a, b))  # ...but still with the same contents!

a[0][1] = 561          # a changes...
print(b)               # ...but b is still all ones
tensor([[True, True],
        [True, True]])
tensor([[1., 1.],
        [1., 1.]])

在使用 ``clone()`` 时有一个重要事项需要注意。 如果您的源张量启用了 autograd,那么克隆的张量也会启用。这将在关于 autograd 的视频中更深入地讨论, 但如果您想要轻量级的详细信息,请继续。

在许多情况下,这将是您想要的。 例如,如果您的模型在其 forward() 方法中有多个计算路径,并且原始张量及其克隆都贡献于模型的输出,那么为了启用模型学习,您希望为两个张量启用 autograd。如果您的源张量启用了 autograd(如果它是学习权重集合或源自涉及权重的计算,通常会如此),那么您将获得想要的结果。

另一方面,如果您正在进行一项计算,其中原始张量和它的克隆都不需要跟踪梯度,那么只要源张量禁用了 autograd,您就可以进行。

还有第三种情况: 想象一下您正在模型 forward() 函数中执行一项计算,其中默认情况下所有项的梯度都已开启,但您想在过程中提取一些值来生成一些指标。在这种情况下,您希望源张量的克隆副本跟踪梯度 - autograd 的历史跟踪关闭可以提高性能。为此,您可以使用源张量上的 .detach() 方法。

a = torch.rand(2, 2, requires_grad=True) # turn on autograd
print(a)

b = a.clone()
print(b)

c = a.detach().clone()
print(c)

print(a)
tensor([[0.0905, 0.4485],
        [0.8740, 0.2526]], requires_grad=True)
tensor([[0.0905, 0.4485],
        [0.8740, 0.2526]], grad_fn=<CloneBackward0>)
tensor([[0.0905, 0.4485],
        [0.8740, 0.2526]])
tensor([[0.0905, 0.4485],
        [0.8740, 0.2526]], requires_grad=True)

这里发生了什么?

  • 我们创建了 a,并将 requires_grad=True 开启。我们还没有介绍这个可选参数,但将在关于 autograd 的单元中介绍。

  • 当我们打印 a 时,它通知我们属性 requires_grad=True - 这意味着 autograd 和计算历史跟踪已开启。

  • 我们克隆了 a 并将其标记为 b。当我们打印 b 时,我们可以看到它正在跟踪其计算历史 - 它继承了 a 的 autograd 设置,并添加到计算历史记录中。

  • 我们将 a 克隆到 c,但在此之前调用了 detach()

  • 打印 c 时,我们看不到计算历史,也没有 requires_grad=True

detach() 方法将张量与计算历史分离。 它表示“像 autograd 关闭一样继续执行接下来的操作。” 它这样做不会改变 a - 当我们在最后再次打印 a 时,您可以看到它保留了 requires_grad=True 属性。

移动到加速器#

PyTorch 的主要优势之一是其在 CUDA、MPS、MTIA 或 XPU 等加速器上的强大加速功能。到目前为止,我们所做的一切都是在 CPU 上完成的。如何移动到更快的硬件?

首先,我们应该使用 is_available() 方法检查加速器是否可用。

注意

如果您没有加速器,本节中的可执行单元将不会执行任何与加速器相关的代码。

if torch.accelerator.is_available():
    print('We have an accelerator!')
else:
    print('Sorry, CPU only.')
We have an accelerator!

一旦我们确定一个或多个加速器可用,我们就需要将数据放在加速器可以看到的地方。您的 CPU 在您计算机的 RAM 中对数据进行计算。您的加速器有专用的内存与之连接。每当您想在设备上执行计算时,都必须将该计算所需的所有数据移动到该设备可访问的内存中。(俗称,“将数据移动到 GPU 可访问内存”简称为“将数据移动到 GPU”。)

有多种方法可以将数据加载到目标设备。您可以在创建时执行此操作。

if torch.accelerator.is_available():
    gpu_rand = torch.rand(2, 2, device=torch.accelerator.current_accelerator())
    print(gpu_rand)
else:
    print('Sorry, CPU only.')
tensor([[0.3344, 0.2640],
        [0.2119, 0.0582]], device='cuda:0')

默认情况下,新张量在 CPU 上创建,因此当我们需要在加速器上创建张量时,必须使用可选的 device 参数进行指定。您可以从打印新张量时看到,PyTorch 会告知我们它在哪一个设备上(如果它不在 CPU 上)。

您可以使用 torch.accelerator.device_count() 查询加速器数量。如果您有多个加速器,则可以通过索引指定它们,以 CUDA 为例:device='cuda:0'device='cuda:1' 等。

作为编码实践,使用字符串常量在所有地方指定我们的设备非常脆弱。在理想情况下,无论是在 CPU 还是加速器硬件上,您的代码都能稳健运行。您可以通过创建一个设备句柄来实现这一点,该句柄可以传递给您的张量而不是字符串。

my_device = torch.accelerator.current_accelerator() if torch.accelerator.is_available() else torch.device('cpu')
print('Device: {}'.format(my_device))

x = torch.rand(2, 2, device=my_device)
print(x)
Device: cuda
tensor([[0.0024, 0.6778],
        [0.2441, 0.6812]], device='cuda:0')

如果您有一个现有张量位于某个设备上,您可以使用 to() 方法将其移动到另一个设备。下面一行代码在 CPU 上创建一个张量,并将其移动到您在上一个单元中获取的任何设备句柄。

y = torch.rand(2, 2)
y = y.to(my_device)

重要的是要知道,为了进行涉及两个或多个张量的计算,所有张量必须位于同一设备上。下面的代码将引发运行时错误,无论您是否有可用的加速器设备,以 CUDA 为例:

x = torch.rand(2, 2)
y = torch.rand(2, 2, device='cuda')
z = x + y  # exception will be thrown

操作张量形状#

有时,您需要更改张量的形状。下面,我们将介绍一些常见情况以及如何处理它们。

更改维数#

您可能需要更改维数的一种情况是向模型传递单个输入实例。PyTorch 模型通常期望输入批次

例如,假设有一个模型处理 3 x 226 x 226 的图像 - 一个 226 像素的正方形,带有 3 个颜色通道。当您加载并转换它时,您将得到一个形状为 (3, 226, 226) 的张量。但是,您的模型期望输入形状为 (N, 3, 226, 226),其中 N 是批次中的图像数量。那么如何创建一个包含一个的批次呢?

a = torch.rand(3, 226, 226)
b = a.unsqueeze(0)

print(a.shape)
print(b.shape)
torch.Size([3, 226, 226])
torch.Size([1, 3, 226, 226])

unsqueeze() 方法添加一个范围为 1 的维度。 unsqueeze(0) 将其添加为新的零维 - 现在您有一个包含一个的批次!

那么,如果这是压缩(unsqueezing)?我们说的压缩(squeezing)是什么意思?我们利用了任何范围为 1 的维度改变张量中元素数量这一事实。

c = torch.rand(1, 1, 1, 1, 1)
print(c)
tensor([[[[[0.2347]]]]])

继续上面的示例,假设模型的输出是每个输入的 20 元素向量。那么您期望输出的形状为 (N, 20),其中 N 是输入批次中的实例数量。这意味着对于我们只有一个输入的批次,我们将得到一个形状为 (1, 20) 的输出。

如果您想使用该输出进行一些非批处理的计算 - 某种只需要 20 元素向量的东西?

a = torch.rand(1, 20)
print(a.shape)
print(a)

b = a.squeeze(0)
print(b.shape)
print(b)

c = torch.rand(2, 2)
print(c.shape)

d = c.squeeze(0)
print(d.shape)
torch.Size([1, 20])
tensor([[0.1899, 0.4067, 0.1519, 0.1506, 0.9585, 0.7756, 0.8973, 0.4929, 0.2367,
         0.8194, 0.4509, 0.2690, 0.8381, 0.8207, 0.6818, 0.5057, 0.9335, 0.9769,
         0.2792, 0.3277]])
torch.Size([20])
tensor([0.1899, 0.4067, 0.1519, 0.1506, 0.9585, 0.7756, 0.8973, 0.4929, 0.2367,
        0.8194, 0.4509, 0.2690, 0.8381, 0.8207, 0.6818, 0.5057, 0.9335, 0.9769,
        0.2792, 0.3277])
torch.Size([2, 2])
torch.Size([2, 2])

您可以从形状中看出,我们的 2 维张量现在是 1 维的,如果您仔细查看上面单元格的输出,您会看到 a 的打印显示了一个“额外的”方括号 [],这是由于有一个额外的维度。

您只能压缩范围为 1 的维度。请参见上面,我们尝试压缩 c 中大小为 2 的维度,但返回了与开始时相同的形状。对 squeeze()unsqueeze() 的调用只能作用于范围为 1 的维度,因为否则会改变张量中的元素数量。

您也可以使用 unsqueeze() 来简化广播。回想一下上面我们有以下代码的示例:

a = torch.ones(4, 3, 2)

c = a * torch.rand(   3, 1) # 3rd dim = 1, 2nd dim identical to a
print(c)

其最终效果是将运算广播到维度 0 和 2,导致随机的 3 x 1 张量与 a 中每个 3 元素列进行逐元素乘法。

如果随机向量只是一个 3 元素向量呢?我们将失去进行广播的能力,因为最后一个维度将不匹配广播规则。 unsqueeze() 来救援。

a = torch.ones(4, 3, 2)
b = torch.rand(   3)     # trying to multiply a * b will give a runtime error
c = b.unsqueeze(1)       # change to a 2-dimensional tensor, adding new dim at the end
print(c.shape)
print(a * c)             # broadcasting works again!
torch.Size([3, 1])
tensor([[[0.1891, 0.1891],
         [0.3952, 0.3952],
         [0.9176, 0.9176]],

        [[0.1891, 0.1891],
         [0.3952, 0.3952],
         [0.9176, 0.9176]],

        [[0.1891, 0.1891],
         [0.3952, 0.3952],
         [0.9176, 0.9176]],

        [[0.1891, 0.1891],
         [0.3952, 0.3952],
         [0.9176, 0.9176]]])

squeeze()unsqueeze() 方法也有原地版本,squeeze_()unsqueeze_()

batch_me = torch.rand(3, 226, 226)
print(batch_me.shape)
batch_me.unsqueeze_(0)
print(batch_me.shape)
torch.Size([3, 226, 226])
torch.Size([1, 3, 226, 226])

有时您会想更彻底地改变张量的形状,同时保留元素数量及其内容。这种情况发生在一个模型中的卷积层和线性层之间的接口处 - 在图像分类模型中很常见。卷积核会产生形状为特征 x 宽度 x 高度的输出张量,但接下来的线性层需要一个 1 维输入。 reshape() 会为您完成此操作,前提是您请求的维度产生的元素数量与输入张量相同。

output3d = torch.rand(6, 20, 20)
print(output3d.shape)

input1d = output3d.reshape(6 * 20 * 20)
print(input1d.shape)

# can also call it as a method on the torch module:
print(torch.reshape(output3d, (6 * 20 * 20,)).shape)
torch.Size([6, 20, 20])
torch.Size([2400])
torch.Size([2400])

注意

单元格最后一行中的 (6 * 20 * 20,) 参数是因为 PyTorch 在指定张量形状时期望一个**元组** - 但当形状是方法的第一个参数时,它允许我们作弊,只使用一系列整数。在这里,我们必须添加括号和逗号,以说服该方法这实际上是一个单元素元组。

reshape() 能够做到时,它会返回一个对要更改的张量的视图 - 即,一个查看相同底层内存区域的独立张量对象。这一点很重要: 这意味着对源张量的任何更改都将反映在该张量的视图上,除非您对其进行 clone()

存在一些超出本入门范围的条件,在这些条件下 reshape() 必须返回一个包含数据副本的张量。有关更多信息,请参阅 文档

NumPy 桥接#

在上面关于广播的部分,我们提到 PyTorch 的广播语义与 NumPy 的兼容 - 但 PyTorch 和 NumPy 之间的联系比这更深。

如果您有现有的 ML 或科学代码,其中数据存储在 NumPy ndarrays 中,您可能希望将相同的数据表示为 PyTorch 张量,无论是为了利用 PyTorch 的 GPU 加速,还是它用于构建 ML 模型的高效抽象。在 ndarrays 和 PyTorch 张量之间进行切换很容易。

import numpy as np

numpy_array = np.ones((2, 3))
print(numpy_array)

pytorch_tensor = torch.from_numpy(numpy_array)
print(pytorch_tensor)
[[1. 1. 1.]
 [1. 1. 1.]]
tensor([[1., 1., 1.],
        [1., 1., 1.]], dtype=torch.float64)

PyTorch 创建一个与 NumPy 数组形状相同且包含相同数据的张量,甚至保留 NumPy 的默认 64 位浮点数据类型。

转换也可以同样轻松地反向进行。

pytorch_rand = torch.rand(2, 3)
print(pytorch_rand)

numpy_rand = pytorch_rand.numpy()
print(numpy_rand)
tensor([[0.8716, 0.2459, 0.3499],
        [0.2853, 0.9091, 0.5695]])
[[0.87163675 0.2458961  0.34993553]
 [0.2853077  0.90905803 0.5695162 ]]

重要的是要了解,这些转换后的对象使用的是与其源对象相同的底层内存,这意味着对其中一个的更改会反映在另一个中。

numpy_array[1, 1] = 23
print(pytorch_tensor)

pytorch_rand[1, 1] = 17
print(numpy_rand)
tensor([[ 1.,  1.,  1.],
        [ 1., 23.,  1.]], dtype=torch.float64)
[[ 0.87163675  0.2458961   0.34993553]
 [ 0.2853077  17.          0.5695162 ]]

脚本总运行时间: (0 分钟 0.443 秒)