Export the cosine_similarity op as an ATenOp correctly (#21884)

houseroad · facebook-github-bot · commit f1c1d1a96416 · 2019-06-17T23:34:59.000-07:00
Summary: cosine_similarity has two non-tensor parameters, needs some special handling. Add the support for its export in this diff. Pull Request resolved: #21884 Reviewed By: zrphercule Differential Revision: D15866807 Pulled By: houseroad fbshipit-source-id: a165fbc00c65c44b276df89ae705ca8960349d48
diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py
@@ -891,6 +891,13 @@ def forward(self, x):
         x = torch.randn(*shape)
         self.run_model_test(MyModel(), train=False, input=(x,), batch_size=BATCH_SIZE, use_gpu=False)
 
+    def test_cosine_similarity(self):
+        shape = (100, 128)
+        x = torch.randn(*shape)
+        y = torch.randn(*shape)
+        self.run_model_test(torch.nn.CosineSimilarity(dim=1, eps=1e-6), train=False,
+                            input=(x, y), batch_size=BATCH_SIZE, use_gpu=False)
+
     def test_lstm_constant_folding(self):
         class LstmNet(nn.Module):
             def __init__(self, input_size, hidden_size, num_layers, bidirectional):
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -958,6 +958,11 @@ def layer_norm(g, self, normalized_shape, weight, bias, eps, cudnn_enable):
                 eps_f=eps, cudnn_enable_i=cudnn_enable, operator_s="layer_norm")
 
 
+@parse_args('v', 'v', 'i', 'f')
+def cosine_similarity(g, x1, x2, dim, eps):
+    return g.op("ATen", x1, x2, dim_i=dim, eps_f=eps, operator_s="cosine_similarity")
+
+
 # ignore clone operators that are inserted by PyTorch autograd
 def clone(g, input):
     return input
