[Torch] Support hard_swish op

python/tvm/relay/frontend/pytorch.py

@@ -790,6 +790,16 @@ def log_sigmoid(self, inputs, input_types):
         data = inputs[0]
         return _op.log(_op.tensor.sigmoid(data))
 
+    def hard_swish(self, inputs, input_types):
+        input = inputs[0]
+        dtype = input_types[0]
+
+        def _relu6(data):
+            return _op.tensor.clip(data, 0.0, 6.0)
+
+        return input * _relu6(input + _expr.const(3.0,  dtype=dtype)) / _expr.const(6.0, dtype=dtype)
+
+
     def adaptive_avg_pool_2d(self, inputs, input_types):
         data = inputs[0]
         output_size = inputs[1]
@@ -2266,6 +2276,8 @@ def create_convert_map(self):
             "aten::bincount": self.bincount,
             "aten::scatter_add": self.scatter_add,
             "aten::__not__": self.logical_not,
+            "aten::hardswish_": self.hard_swish,
+            "aten::hardswish": self.hard_swish,
         }
 
     def update_convert_map(self, custom_map):

tests/python/frontend/pytorch/test_forward.py

@@ -3437,6 +3437,16 @@ def test_fn(x, weights=None):
     verify_trace_model(test_fn, [inp, weights], targets)
 
 
+def test_hard_swish():
+    hard_swish = torch.nn.Hardswish()
+    hard_swish_inplace = torch.nn.Hardswish(inplace=True)
+    examples = [torch.rand(8), torch.rand(8, 8), torch.rand(8, 16), torch.rand(1, 1, 8)]
+    targets = ["llvm", "cuda"]
+    for input in examples:
+        verify_trace_model(hard_swish, [input], targets)
+        verify_trace_model(hard_swish_inplace, [input], targets)
+
+
 if __name__ == "__main__":
     # some structural tests
     test_forward_traced_function()
@@ -3603,3 +3613,4 @@ def test_fn(x, weights=None):
 
     # Test convert torch script(jit) with specific inputs' types
     test_convert_torch_script_with_input_types()
+    test_hard_swish()