PyTorch中的Eager模式

NLP工程化2024-01-18 18:13

在 PyTorch 中，"Eager Execution"（即即时执行模式）是指一种动态图计算模式，其中每个操作都立即被执行，而不是被先放入计算图中。这与静态图计算框架（如 TensorFlow 的早期版本）的工作方式不同。在即时执行模式中，你可以像使用 NumPy 一样进行操作，逐步构建计算图，方便调试和交互。

一. Eager Execution 基本概念

1. 即时计算

在 Eager Execution 模式下，每个操作都会立即执行，而不是构建一个计算图。

2. 调试友好

Eager Execution 模式更加直观，更容易进行调试，因为可以像在传统编程中一样逐行执行代码。

3. 动态构建计算图

在 Eager 模式下，计算图是动态构建的，这意味着可以使用控制流结构（如 if 语句、循环等）而无需事先定义静态图。

二. Eager Execution 示例

在这个示例中，首先创建了两个张量 x 和 y，然后进行了一些操作，最后计算了它们的平均值。在 Eager Execution 模式下，每个操作都会立即执行，而不需要构建静态图。最后，我们使用反向传播计算了梯度。

复制代码

import torch

# 启用 Eager Execution 模式
torch.set_grad_enabled(True)  # 默认情况下是开启的，这里仅为演示

# 创建两个张量
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)

# 执行操作
z = x + y
result = z.mean()

# 输出结果
print("x:", x)
print("y:", y)
print("z:", z)
print("result:", result)

# 反向传播
result.backward()

# 打印梯度
print("Gradient of x:", x.grad)
print("Gradient of y:", y.grad)

请注意，Eager Execution 模式在 PyTorch 1.0 版本之后成为默认行为，无需额外的设置。