PyTorch torch.compile 関数
Pytorch torch リファレンスマニュアルtorch.compilePyTorchでモデルをコンパイルして最適化するための関数であり、モデルの実行性能を大幅に向上させることができます。
関数定義
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)
その他の拡張