注意
点击 这里 下载完整示例代码
简介 || 张量 || 自动求导 || 构建模型 || TensorBoard 支持 || 训练模型 || 模型理解
PyTorch 张量简介¶
创建时间:2021年11月30日 | 最后更新:2024年8月1日 | 最后验证:2024年11月5日
请跟随下方视频或在 YouTube 上观看。
张量是PyTorch中的核心数据抽象。此交互式笔记本提供了对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.9601e+10, 3.0637e-41, 0.0000e+00, 2.3510e-38],
[ 1.5736e-27, 3.0630e-41, 1.0842e-19, 0.0000e+00],
[ 1.6411e-27, 3.0630e-41, 1.5773e-27, 3.0630e-41]])
我们刚刚做了什么可以这样拆解:
我们使用了
torch模块中的众多工厂方法之一创建了一个张量。张量本身是二维的,有3行和4列。
对象返回的类型是
torch.Tensor,它是torch.FloatTensor的别名;默认情况下,PyTorch 张量填充的是 32 位浮点数。(更多关于数据类型的信息请参见下方。)打印你的张量时,可能会看到一些随机值。
torch.empty()调用为张量分配了内存,但没有初始化其值 - 因此你看到的是内存分配时的原有数据。
关于张量及其维度和相关术语的简要说明:
您有时会看到一个一维张量称为 向量。
likewise, 一个二维张量通常被称为 矩阵。
Anything 有超过两个维度的话,通常就叫做张量。
More often than not, 你通常会希望初始化你的张量为某些值。常见的案例是全零、全一或随机值,而 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]])
所有的工厂方法都像你预期的那样工作——我们有一个全是零的张量,另一个全是one的张量,还有一个包含0到1之间随机值的张量。
随机张量和种子设置¶
讨论随机张量时,您注意到它前面立即调用了
torch.manual_seed() 吗?用随机值初始化张量,例如模型的学习权重,是很常见的做法,但在研究环境中,您可能希望确保结果的可重复性。手动设置随机数生成器的种子就是实现这一目标的方法。让我们更详细地了解一下:
torch.manual_seed(1729)
random1 = torch.rand(2, 3)
print(random1)
random2 = torch.rand(2, 3)
print(random2)
torch.manual_seed(1729)
random3 = torch.rand(2, 3)
print(random3)
random4 = torch.rand(2, 3)
print(random4)
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]])
您应该看到上面的是 random1 和 random3 携带相同的值,同样地,random2 和 random4 也携带相同的值。手动设置随机数生成器(RNG)的种子会重置它,因此大多数情况下,依赖随机数的相同计算应提供相同的結果。
有关更多信息,请参阅PyTorch 关于可重复性的文档。
张量形状¶
Often, 当你在对两个或多个张量进行操作时,它们需要具有相同的 形状 - 即,具有相同数量的维度和每个维度中相同数量的单元格。为此,我们有 torch.*_like() 种方法:
x = torch.empty(2, 2, 3)
print(x.shape)
print(x)
empty_like_x = torch.empty_like(x)
print(empty_like_x.shape)
print(empty_like_x)
zeros_like_x = torch.zeros_like(x)
print(zeros_like_x.shape)
print(zeros_like_x)
ones_like_x = torch.ones_like(x)
print(ones_like_x.shape)
print(ones_like_x)
rand_like_x = torch.rand_like(x)
print(rand_like_x.shape)
print(rand_like_x)
torch.Size([2, 2, 3])
tensor([[[-7.0046e-35, 3.0630e-41, -7.0202e-35],
[ 3.0630e-41, 8.9683e-44, 0.0000e+00]],
[[ 1.1210e-43, 0.0000e+00, -9.1574e+14],
[ 3.0637e-41, -7.0381e-35, 3.0630e-41]]])
torch.Size([2, 2, 3])
tensor([[[-9.4520e+10, 3.0637e-41, 1.4013e-45],
[ 0.0000e+00, 1.4013e-45, 0.0000e+00]],
[[ 1.4013e-45, 0.0000e+00, 1.4013e-45],
[ 0.0000e+00, 1.4013e-45, 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]]])
The first new thing in the code cell above is the use of the .shape
属性 on a tensor. This property contains a list of the extent of
each dimension of a tensor - in our case, x 是一个三维张量,形状为 2 x 2 x 3。
Below that,我们调用 .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)
创建张量时设置其基础数据类型的最简单方法是在创建时使用可选参数。在上面单元格的第一行中,我们将dtype=torch.int16设置为张量a的值。当我们打印a时,我们可以看到它充满了1而不是1.——这是Python提示这是一个整数类型而非浮点数类型的微妙线索。
另一个需要注意的是打印 a 时,与我们将 dtype 留作默认值(32位浮点数)时不同,打印张量还会指定其 dtype。
您可能还注意到了,我们从使用一系列整数参数来指定张量的形状,转变为将这些参数组合在一个元组中。这并不是必要的——PyTorch 会将一系列初始且未标注的整数参数视为张量的形状——但在添加可选参数时,这样做可以使您的意图更加清晰易读。
The other way to set the datatype is with the .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的那行代码中所做的那样。
两个张量之间的相似操作也像你预期的那样进行:
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 ndarrays中的广播语义,您会发现在这里同样适用这些规则。
同形状规则的例外情况是张量广播。这里有一个例子:
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 (其中 a 和 b 是张量),新的张量 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])
对于算术运算,有类似的行为函数:
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中的任何对象一样,将张量赋值给一个变量会使该变量成为张量的标签,并不会复制它。例如:
tensor([[ 1., 561.],
[ 1., 1.]])
但如果你想要一份独立的数据副本进行操作呢?clone() 方法正是为你准备的:
tensor([[True, True],
[True, True]])
tensor([[1., 1.],
[1., 1.]])
使用 ``clone()`` 时需要注意一个重要事项。 如果您的源张量启用了自动求导,则克隆的张量也会启用自动求导。 这将在自动求导的视频中更深入地讨论, 但如果您想了解简化版的细节,请继续阅读。
在许多情况下,这正是你想要的。例如,如果你的模型在forward()方法中有多个计算路径,并且原始张量及其克隆都对模型的输出有所贡献,那么为了使模型能够学习,你需要同时开启两个张量的自动求导功能。如果源张量的自动求导功能已开启(通常情况下,如果它是学习权重的一组或者是从涉及权重的计算中派生出来的),那么你就能得到想要的结果。
另一方面,如果您的计算中既不需要跟踪梯度的原始张量也不需要其克隆,则只要源张量的自动求导已关闭,您就可以继续进行。
虽然还有一种情况:想象你在模型的forward()函数中进行计算,梯度默认情况下对所有内容都是开启状态,但你想中途提取一些值来生成一些指标。在这种情况下,你不希望源张量的克隆副本跟踪梯度——通过关闭autograd的历史记录追踪可以提高性能。为此,你可以使用源张量上的.detach()方法:
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。 我们还没有覆盖这个可选参数,但会在自动求导单元中讲解。当我们打印
a,它会告诉我们属性requires_grad=True- 这意味着自动求导和计算历史跟踪已开启。我们克隆
a并将其标记为b。当我们打印b时,可以看到它正在跟踪其计算历史 - 它继承了a的自动求导设置,并添加到了计算历史中。我们把
a克隆到c,但我们会先调用detach()。Printing
c,我们看到没有计算历史记录,也没有requires_grad=True。
The detach() 方法 断开了张量与其计算历史的连接。 它表示,“接下来的操作就好像自动求导被关闭了一样。” 这个操作 不会改变 a - 你可以看到当我们再次打印 a 时,它仍然保留了 requires_grad=True 属性。
移动到GPU¶
PyTorch的一个主要优势是其在Nvidia兼容CUDA的GPU上的强大加速。(“CUDA”代表计算统一设备架构,这是Nvidia的并行计算平台。)到目前为止,我们所做的所有操作都在CPU上进行。我们如何切换到更快的硬件?
首先,我们应该使用
is_available() 方法检查是否可用 GPU。
注意
如果没有安装CUDA兼容的GPU和CUDA驱动程序,本节中的可执行代码单元将不会运行任何与GPU相关的代码。
if torch.cuda.is_available():
print('We have a GPU!')
else:
print('Sorry, CPU only.')
We have a GPU!
一旦确定了一台或多台GPU可用,我们需要将数据放在GPU可以访问的地方。你的CPU在计算机的RAM中进行计算。你的GPU则有专门的内存。每当你要在某个设备上执行计算时,必须将该计算所需的所有数据移动到该设备可访问的内存中。(通俗地说,“将数据移动到GPU可访问的内存中”通常简称作“将数据移动到GPU”。)
有多种方法可以将数据传输到目标设备。您可能在创建时进行数据传输:
if torch.cuda.is_available():
gpu_rand = torch.rand(2, 2, device='cuda')
print(gpu_rand)
else:
print('Sorry, CPU only.')
tensor([[0.3344, 0.2640],
[0.2119, 0.0582]], device='cuda:0')
默认情况下,新张量在CPU上创建,所以我们必须指定使用可选的
device 参数时要在GPU上创建张量。当我们打印新张量时,PyTorch 会告诉我们它位于哪个设备(如果不是CPU)。
您可以查询GPU数量:torch.cuda.device_count()。如果您有多个GPU,可以通过索引指定它们:
device='cuda:0',device='cuda:1',等等。
作为编码练习,使用字符串常量来指定设备在各个方面都比较脆弱。在理想的世界里,无论你是使用CPU还是GPU硬件,你的代码都应该能够稳健运行。你可以通过创建一个设备句柄并将其传递给张量而不是字符串来实现这一点:
if torch.cuda.is_available():
my_device = torch.device('cuda')
else:
my_device = torch.device('cpu')
print('Device: {}'.format(my_device))
x =