[quant][graph] Graph mode quantization support for sigmoid (#36622)

supriyar · facebook-github-bot · commit 753157b88eeb · 2020-04-16T16:50:54.000-07:00
Summary: Pull Request resolved: #36622 Test Plan: python test/quantization/test_quantize_script.py test_swap_dequantize_all_ops Imported from OSS Differential Revision: D21075255 fbshipit-source-id: 025f432215eaa8acf34d492e7722102ca053abeb
diff --git a/test/quantization/test_quantize_script.py b/test/quantization/test_quantize_script.py
@@ -1459,6 +1459,7 @@ def __init__(self):
                 self.avgpool2d = torch.nn.AvgPool2d(3)
                 self.avgpool3d = torch.nn.AvgPool3d(3)
                 self.conv = torch.nn.Conv2d(3, 3, 3)
+                self.sigmoid = torch.nn.Sigmoid()
 
             def forward(self, x):
                 x = self.conv(x)
@@ -1481,7 +1482,9 @@ def forward(self, x):
                 x = torch.max(x)
                 x = torch.min(x)
                 x = torch.mean(x)
+                x = torch.sigmoid(x)
                 x = x.reshape([-1])
+                x = self.sigmoid(x)
                 x = x.view(-1)
                 x = x.transpose(1, 2)
                 x = x.contiguous()
@@ -1493,6 +1496,7 @@ def forward(self, x):
                 x = F.upsample(x, (32, 32))
                 x = F.upsample_bilinear(x, (32, 32))
                 x = F.upsample_nearest(x, (32, 32))
+                x = F.sigmoid(x)
                 x = x.permute(0, 2, 3, 1)
                 x = torch.repeat_interleave(x, 3, 1)
                 x = self.conv(x)
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -86,6 +86,7 @@ std::vector<std::string> _single_input_general_call_funcs = {
     "upsample_bilinear",
     "upsample_nearest",
     "relu",
+    "sigmoid",
 };
 
 // Similar to prim::CallFunctions, there are aten ops that doesn't
@@ -121,6 +122,7 @@ std::vector<std::string> _single_input_general_aten_funcs = {
     "permute",
     "repeat_interleave",
     "relu",
+    "sigmoid",
 };
 
 struct FuncArg {
