TorchScript 中的模型凍結

創建於：2020 年 7 月 28 日 | 最後更新於：2024 年 12 月 2 日 | 最後驗證於：2024 年 11 月 5 日

警告

TorchScript 已不再積極開發中。

在本教程中，我們將介紹 TorchScript 中模型凍結的語法。凍結是將 PyTorch 模組引數和屬性值內聯到 TorchScript 內部表示中的過程。引數和屬性值被視為最終值，在生成的凍結模組中無法修改它們。

基本語法

模型凍結可以使用以下 API 呼叫

torch.jit.freeze(mod : ScriptModule, names : str[]) -> ScriptModule

注意，輸入的模組可以是指令碼化或跟蹤的結果。請參閱 https://pytorch.com.tw/tutorials/beginner/Intro_to_TorchScript_tutorial.html

接下來，我們透過一個示例演示凍結的工作原理

import torch, time

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, 3, 1)
        self.conv2 = torch.nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = torch.nn.Dropout2d(0.25)
        self.dropout2 = torch.nn.Dropout2d(0.5)
        self.fc1 = torch.nn.Linear(9216, 128)
        self.fc2 = torch.nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = torch.nn.functional.relu(x)
        x = self.conv2(x)
        x = torch.nn.functional.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = torch.nn.functional.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = torch.nn.functional.log_softmax(x, dim=1)
        return output

    @torch.jit.export
    def version(self):
        return 1.0

net = torch.jit.script(Net())
fnet = torch.jit.freeze(net)

print(net.conv1.weight.size())
print(net.conv1.bias)

try:
    print(fnet.conv1.bias)
    # without exception handling, prints:
    # RuntimeError: __torch__.z.___torch_mangle_3.Net does not have a field
    # with name 'conv1'
except RuntimeError:
    print("field 'conv1' is inlined. It does not exist in 'fnet'")

try:
    fnet.version()
    # without exception handling, prints:
    # RuntimeError: __torch__.z.___torch_mangle_3.Net does not have a field
    # with name 'version'
except RuntimeError:
    print("method 'version' is not deleted in fnet. Only 'forward' is preserved")

fnet2 = torch.jit.freeze(net, ["version"])

print(fnet2.version())

B=1
warmup = 1
iter = 1000
input = torch.rand(B, 1,28, 28)

start = time.time()
for i in range(warmup):
    net(input)
end = time.time()
print("Scripted - Warm up time: {0:7.4f}".format(end-start), flush=True)

start = time.time()
for i in range(warmup):
    fnet(input)
end = time.time()
print("Frozen   - Warm up time: {0:7.4f}".format(end-start), flush=True)

start = time.time()
for i in range(iter):
    input = torch.rand(B, 1,28, 28)
    net(input)
end = time.time()
print("Scripted - Inference: {0:5.2f}".format(end-start), flush=True)

start = time.time()
for i in range(iter):
    input = torch.rand(B, 1,28, 28)
    fnet2(input)
end = time.time()
print("Frozen    - Inference time: {0:5.2f}".format(end-start), flush =True)

在我的機器上，我測量了時間

• 指令碼化模型 - 熱啟動時間: 0.0107

• 凍結模型 - 熱啟動時間: 0.0048

• 指令碼化模型 - 推理時間: 1.35

• 凍結模型 - 推理時間: 1.17

在我們的示例中，熱啟動時間衡量了前兩次執行的時間。凍結模型比指令碼化模型快 50%。在某些更復雜的模型上，我們觀察到熱啟動時間有更高的加速效果。凍結之所以能實現這種加速，是因為它執行了 TorchScript 在前幾次執行時必須做的一些工作。

推理時間衡量模型熱啟動後的推理執行時間。儘管我們在執行時間上觀察到顯著差異，但凍結模型通常比指令碼化模型快約 15%。當輸入較大時，加速效果較小，因為執行時間主要由張量操作決定。

結論

在本教程中，我們學習了模型凍結。凍結是一種最佳化模型以進行推理的有用技術，它還可以顯著減少 TorchScript 的熱啟動時間。