-
Notifications
You must be signed in to change notification settings - Fork 214
/
activations_autofn.py
72 lines (52 loc) · 1.92 KB
/
activations_autofn.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import torch
from torch import nn as nn
from torch.nn import functional as F
__all__ = ['swish_auto', 'SwishAuto', 'mish_auto', 'MishAuto']
class SwishAutoFn(torch.autograd.Function):
"""Swish - Described in: https://arxiv.org/abs/1710.05941
Memory efficient variant from:
https://medium.com/the-artificial-impostor/more-memory-efficient-swish-activation-function-e07c22c12a76
"""
@staticmethod
def forward(ctx, x):
result = x.mul(torch.sigmoid(x))
ctx.save_for_backward(x)
return result
@staticmethod
def backward(ctx, grad_output):
x = ctx.saved_tensors[0]
x_sigmoid = torch.sigmoid(x)
return grad_output.mul(x_sigmoid * (1 + x * (1 - x_sigmoid)))
def swish_auto(x, inplace=False):
# inplace ignored
return SwishAutoFn.apply(x)
class SwishAuto(nn.Module):
def __init__(self, inplace: bool = False):
super(SwishAuto, self).__init__()
self.inplace = inplace
def forward(self, x):
return SwishAutoFn.apply(x)
class MishAutoFn(torch.autograd.Function):
"""Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
Experimental memory-efficient variant
"""
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
y = x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x)))
return y
@staticmethod
def backward(ctx, grad_output):
x = ctx.saved_tensors[0]
x_sigmoid = torch.sigmoid(x)
x_tanh_sp = F.softplus(x).tanh()
return grad_output.mul(x_tanh_sp + x * x_sigmoid * (1 - x_tanh_sp * x_tanh_sp))
def mish_auto(x, inplace=False):
# inplace ignored
return MishAutoFn.apply(x)
class MishAuto(nn.Module):
def __init__(self, inplace: bool = False):
super(MishAuto, self).__init__()
self.inplace = inplace
def forward(self, x):
return MishAutoFn.apply(x)