在PyTorch中,可以使用requires_grad=True
来标记一个张量需要求导。然后,可以使用backward()
方法来计算该张量的梯度。
例如,假设有一个张量x,我们想要计算它的平方的梯度:
import torch
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
y.backward()
print(x.grad)
输出结果为:
tensor([2., 4., 6.])
这表示x的梯度为[2, 4, 6]。
另外,还可以使用autograd
自动计算梯度。只需将输入张量的requires_grad
属性设置为True
,PyTorch将自动计算所有需要求导的张量的梯度。例如:
import torch
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
z = y * 2
torch.autograd.backward(z)
print(x.grad)
输出结果同样为:
tensor([2., 4., 6.])