PyTorch torch.compile 函数
Pytorch torch 参考手册torch.compile 是 PyTorch 中用于编译优化模型的函数，可以显著提升模型的运行性能。
函数定义
torch.compile(model, *, mode="default", backend="inductor", fullgraph=False, dynamic=None, options=None, disable=False)
使用示例
实例
import torch
# 定义一个简单的模型
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 10)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# 编译模型
compiled_model = torch.compile(model)
# 使用编译后的模型进行推理
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
# 定义一个简单的模型
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 10)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# 编译模型
compiled_model = torch.compile(model)
# 使用编译后的模型进行推理
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)