mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Merge pull request #401 from hwangdeyu/deyu/add_HardSwishJitAutoFn_operator

Browse Source 
add HardSwishJitAutoFn operator export to onnx
This commit is contained in:
Ross Wightman 2021-02-04 12:18:19 -08:00 committed by GitHub
commit ea36a78cff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
timm/models/layers/activations_me.py
View File

@ -152,6 +152,13 @@ class HardSwishJitAutoFn(torch.autograd.Function):
x = ctx.saved_tensors[0] x = ctx.saved_tensors[0]
return hard_swish_jit_bwd(x, grad_output) return hard_swish_jit_bwd(x, grad_output)
@staticmethod
def symbolic(g, self):
input = g.op("Add", self, g.op('Constant', value_t=torch.tensor(3, dtype=torch.float)))
hardtanh_ = g.op("Clip", input, g.op('Constant', value_t=torch.tensor(0, dtype=torch.float)), g.op('Constant', value_t=torch.tensor(6, dtype=torch.float)))
hardtanh_ = g.op("Div", hardtanh_, g.op('Constant', value_t=torch.tensor(6, dtype=torch.float)))
return g.op("Mul", self, hardtanh_)
def hard_swish_me(x, inplace=False): def hard_swish_me(x, inplace=False):
return HardSwishJitAutoFn.apply(x) return HardSwishJitAutoFn.apply(x)