在PyTorch中,可以通过torch.autograd
模块来进行梯度计算和反向传播。以下是一个简单的示例代码:
import torch
# 创建一个张量并设置requires_grad=True来追踪其计算历史
x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)
# 定义一个计算图
z = x * y
# 计算z相对于x和y的梯度
z.backward()
# 访问x和y的梯度
print(x.grad) # 输出tensor([3.])
print(y.grad) # 输出tensor([2.])
在这个示例中,我们首先创建了两个张量x
和y
,并设置requires_grad=True
来追踪它们的计算历史。然后我们定义了一个计算图z = x * y
,然后调用z.backward()
来计算z
相对于x
和y
的梯度。最后我们可以通过访问x.grad
和y.grad
来获取梯度值。