update optimize_for_inference docs (#64428)

Elias Ellison · facebook-github-bot · commit 7031fbdc63b0 · 2021-09-02T22:17:58.000-07:00
Summary: Pull Request resolved: #64428 Test Plan: Imported from OSS Reviewed By: pbelevich Differential Revision: D30740898 Pulled By: eellison fbshipit-source-id: b94d2c3deb661a6ba048f19e8c1d5e1799667eeb
diff --git a/docs/source/jit.rst b/docs/source/jit.rst
@@ -60,6 +60,7 @@ Creating TorchScript Code
     ScriptModule
     ScriptFunction
     freeze
+    optimize_for_inference
     save
     load
     ignore
diff --git a/torch/jit/_freeze.py b/torch/jit/_freeze.py
@@ -179,6 +179,18 @@ def optimize_for_inference(mod: ScriptModule) -> ScriptModule:
     This is still in prototype, and may have the potential to slow down your model.
     Primary use cases that have been targeted so far have been vision models on cpu
     and gpu to a lesser extent.
+
+    Example (optimizing a module with Conv->Batchnorm)::
+
+        import torch
+        in_channels, out_channels = 3, 32
+        conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=True)
+        bn = torch.nn.BatchNorm2d(out_channels, eps=.001)
+        mod = torch.nn.Sequential(conv, bn)
+        frozen_mod = torch.jit.optimize_for_inference(torch.jit.script(mod.eval()))
+        assert "batch_norm" not in str(frozen_mod.graph)
+        # if built with MKLDNN, convolution will be run with MKLDNN weights
+        assert "MKLDNN" in frozen_mod.graph
     """
     if not isinstance(mod, ScriptModule):
         raise RuntimeError(
