(测试版) 使用torch.compile编译优化器¶
Created On: Jan 24, 2024 | Last Updated: Jan 29, 2024 | Last Verified: Nov 05, 2024
优化器是训练任何深度学习模型的关键算法。由于它负责更新每个模型参数，因此在大型模型的训练性能中，这常常会成为瓶颈。在此示例中，我们将对优化器应用``torch.compile``以观察其在GPU上的性能改进。
备注
本教程需要PyTorch 2.2.0或更高版本。
模型设置¶
在此示例中，我们将使用一系列简单的线性层。由于我们仅对优化器进行基准测试，模型的选择并不重要，因为优化器性能取决于参数的数量。
取决于您使用的机器，确切结果可能会有所不同。
import torch
model = torch.nn.Sequential(
*[torch.nn.Linear(1024, 1024, False, device="cuda") for _ in range(10)]
)
input = torch.rand(1024, device="cuda")
output = model(input)
output.sum().backward()
设置并运行优化器基准测试¶
在此示例中，我们将使用Adam优化器并创建一个辅助函数将``step()``包装在``torch.compile()``中。
备注
``torch.compile``仅支持计算能力>=7.0的CUDA设备。
# exit cleanly if we are on a device that doesn't support torch.compile
if torch.cuda.get_device_capability() < (7, 0):
print("Exiting because torch.compile is not supported on this device.")
import sys
sys.exit(0)
opt = torch.optim.Adam(model.parameters(), lr=0.01)
@torch.compile(fullgraph=False)
def fn():
opt.step()
# Let's define a helpful benchmarking function:
import torch.utils.benchmark as benchmark
def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
t0 = benchmark.Timer(
stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
)
return t0.blocked_autorange().mean * 1e6
# Warmup runs to compile the function
for _ in range(5):
fn()
eager_runtime = benchmark_torch_function_in_microseconds(opt.step)
compiled_runtime = benchmark_torch_function_in_microseconds(fn)
assert eager_runtime > compiled_runtime
print(f"eager runtime: {eager_runtime}us")
print(f"compiled runtime: {compiled_runtime}us")
示例结果：
Eager运行时：747.2437149845064微秒
编译运行时：392.07384741178微秒
