Add Model Freezing in TorchScript

Author: bzinodev

recipes_source/recipes/torchscript_freezing.py

@@ -0,0 +1,135 @@
+"""
+Model Freezing in TorchScript
+=============================
+
+In this tutorial, we introduce the syntax for *model freezing* in TorchScript.
+Freezing is the process of inlining Pytorch module parameters and attributes
+values into the TorchScript internal representation. Parameter and attribute
+values are treated as final values and they cannot be modified in the resulting
+Frozen module.
+
+Basic Syntax
+------------
+
+Model freezing can be invoked using API below:
+
+* ``torch.jit.freeze(mod : ScriptModule, names : str[]) -> SciptModule``
+
+Note the input module can either be the result of scripting or tracing,
+See https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html
+
+Next, we demonstrate how freezing works using an example:
+"""
+
+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)
+
+"""
+On my machine, I measured the time:
+
+Scripted - Warm up time:  0.0107
+Frozen   - Warm up time:  0.0048
+Scripted - Inference:  1.35
+Frozen   - Inference time:  1.17
+
+In our example, warm up time measures the first two runs. The frozen model
+is 50% faster than the scripted model. On some more complex models, we
+observed even higher speed up of warm up time. freezing achieves this speed up
+because it is doing some the work TorchScript has to do when the first couple
+runs are initiated.
+
+Inference time measures inference execution time after the model is warmed up.
+Although we observed significant variation in execution time, the
+frozen model is often about 15% faster than the scripted model. When input is larger,
+we observe a smaller speed up because the execution is dominated by tensor operations.
+
+Conclusion
+-----------
+
+In this tutorial, we learned about model freezing. Freezing is a useful technique to
+optimize models for inference and it also can significantly reduce TorchScript warmup time.
+"""

recipes_source/recipes_index.rst

@@ -146,6 +146,12 @@ Recipes are bite-sized bite-sized, actionable examples of how to use specific Py
    :link: ../recipes/deployment_with_flask.html
    :tags: Production,TorchScript
 
+.. customcarditem::
+   :header: Model Freezing in TorchScript
+   :card_description: Learn how to use freezing API to optimize your trained model in TorchScript and do inference.
+   :image: ../_static/img/thumbnails/cropped/torchscript_overview.png
+   :link: ../recipes/torchscript_freezing.html
+   :tags: TorchScript
 
 .. End of tutorial card section
 
@@ -180,3 +186,4 @@ Recipes are bite-sized bite-sized, actionable examples of how to use specific Py
    /recipes/recipes/dynamic_quantization
    /recipes/torchscript_inference
    /recipes/deployment_with_flask
+   /recipes/torchscript_freezing