Add mobile_optimized tag to optimized model. (#45479)

Akshit Khurana · facebook-github-bot · commit 5f49d14be261 · 2020-09-29T10:06:57.000-07:00
Summary: Pull Request resolved: #45479 Add a top level boolean attribute to the model called mobile_optimized that is set to true if it is optimized. Test Plan: buck test //caffe2/test:mobile passes Reviewed By: kimishpatel Differential Revision: D23956728 fbshipit-source-id: 79c5931702208b871454319ca2ab8633596b1eb8
diff --git a/test/test_mobile_optimizer.py b/test/test_mobile_optimizer.py
@@ -131,6 +131,23 @@ def forward(self, x):
         bn_input = torch.rand(1, 1, 6, 6)
         torch.testing.assert_allclose(bn_scripted_module(bn_input), no_bn_fold_scripted_module(bn_input), rtol=1e-2, atol=1e-3)
 
+        class MyMobileOptimizedTagTest(torch.nn.Module):
+            def __init__(self):
+                super(MyMobileOptimizedTagTest, self).__init__()
+                self.linear_weight = torch.nn.Parameter(torch.Tensor(torch.rand(linear_weight_shape)))
+                self.linear_bias = torch.nn.Parameter(torch.Tensor(torch.rand((weight_output_dim))))
+
+            def forward(self, x):
+                o = F.linear(x, self.linear_weight, self.linear_bias)
+                return F.relu(o)
+
+        mobile_optimized_tag_module = MyMobileOptimizedTagTest()
+        m = torch.jit.script(mobile_optimized_tag_module)
+        m.eval()
+        opt_m = optimize_for_mobile(m)
+        tag = getattr(opt_m, "mobile_optimized", None)
+        self.assertTrue(tag)
+
         class MyPreserveMethodsTest(torch.nn.Module):
             def __init__(self):
                 super(MyPreserveMethodsTest, self).__init__()
diff --git a/torch/csrc/jit/passes/xnnpack_rewrite.cpp b/torch/csrc/jit/passes/xnnpack_rewrite.cpp
@@ -390,7 +390,7 @@ script::Module optimizeForMobile(
   if (!optimization_blocklist.count(MobileOptimizerType::FUSE_ADD_RELU)) {
     FuseAddRelu(cloned_module);
   }
-
+  cloned_module.register_attribute("mobile_optimized", BoolType::get(), true);
   return cloned_module;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -390,7 +390,7 @@ script::Module optimizeForMobile(`
`390`	`390`	`if (!optimization_blocklist.count(MobileOptimizerType::FUSE_ADD_RELU)) {`
`391`	`391`	`FuseAddRelu(cloned_module);`
`392`	`392`	`}`
`393`		`-`
	`393`	`+ cloned_module.register_attribute("mobile_optimized", BoolType::get(), true);`
`394`	`394`	`return cloned_module;`
`395`	`395`	`}`
`396`	`396`