Export torch.masked_fill with onnx::where

BowenBao · facebook-github-bot · commit 52de34062912 · 2019-07-16T10:55:30.000-07:00
Summary: Pull Request resolved: #22521 Reviewed By: zrphercule Differential Revision: D16155168 Pulled By: houseroad fbshipit-source-id: 5d419f08213324d474b839ba1ae13c799aeee92a
diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py
@@ -1956,6 +1956,24 @@ def forward(self, x):
         inputs = torch.randn(3, 2, 1)
         self.run_model_test(model, train=False, input=(inputs, ), batch_size=BATCH_SIZE)
 
+
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_masked_fill(self):
+        class MaskedFillModel(torch.nn.Module):
+            def forward(self, x):
+                mask = torch.tensor([[0, 0, 1], [1, 1, 0]], dtype=torch.uint8)
+                return x.masked_fill(mask, 2)
+
+        x = torch.zeros(4, 2, 3, requires_grad=True)
+        self.run_model_test(MaskedFillModel(), input=(x, ), train=False, batch_size=BATCH_SIZE)
+
+        class MaskedFillModel2(torch.nn.Module):
+            def forward(self, x):
+                return x.masked_fill(x > 3, -1)
+
+        x = torch.arange(16).view(2, 2, 4).to(torch.float32)
+        self.run_model_test(MaskedFillModel2(), input=(x, ), train=False, batch_size=BATCH_SIZE)
+
 # a bit of metaprogramming to set up all the rnn tests
 
 
diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
@@ -364,6 +364,23 @@ def forward(self, x):
         x = torch.randn(2, 3, 4)
         self.run_test(TensorFactory(), x)
 
+    @skipIfUnsupportedMinOpsetVersion(9)
+    def test_masked_fill(self):
+        class MaskedFillModel(torch.nn.Module):
+            def forward(self, x):
+                mask = torch.tensor([[0, 0, 1], [1, 1, 0]], dtype=torch.uint8)
+                return x.masked_fill(mask, 2)
+
+        x = torch.zeros(4, 2, 3, requires_grad=True)
+        self.run_test(MaskedFillModel(), x)
+
+        class MaskedFillModel2(torch.nn.Module):
+            def forward(self, x):
+                return x.masked_fill(x > 3, -1)
+
+        x = torch.arange(16).view(2, 2, 4).to(torch.float32)
+        self.run_test(MaskedFillModel2(), x)
+
 
 # opset 7 tests
 TestONNXRuntime_opset7 = type(str("TestONNXRuntime_opset7"),
@@ -375,6 +392,7 @@ def forward(self, x):
                               (unittest.TestCase,),
                               dict(TestONNXRuntime.__dict__, opset_version=8))
 
+
 # opset 10 tests
 TestONNXRuntime_opset10 = type(str("TestONNXRuntime_opset10"),
                                (unittest.TestCase,),
diff --git a/torch/onnx/symbolic_opset8.py b/torch/onnx/symbolic_opset8.py
@@ -39,6 +39,7 @@
 
 black_listed_operators = [
     "nonzero", "where", "scatter", "scatter_add", "erf", "sign", "isnan", "gather",
+    "masked_fill"
 ]
 
 for black_listed_op in black_listed_operators:
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -1664,3 +1664,9 @@ def gather(g, self, dim, index, sparse_grad=False):
 @parse_args('v', 'is', 'i')
 def logsumexp(g, input, dim, keepdim):
     return g.op('ReduceLogSumExp', input, axes_i=dim, keepdims_i=keepdim)
+
+
+def masked_fill(g, self, mask, value):
+    mask = _cast_Bool(g, mask, False)
+    value = sym_help._maybe_get_scalar(value)
+    return g.op('Where', mask, sym_help._if_scalar_type_as(g, value, self), self)

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@`
`39`	`39`
`40`	`40`	`black_listed_operators = [`
`41`	`41`	`"nonzero", "where", "scatter", "scatter_add", "erf", "sign", "isnan", "gather",`
	`42`	`+ "masked_fill"`
`42`	`43`	`]`
`43`	`44`
`44`	`45`	`for black_listed_op in black_listed_operators:`