| import mesh_tensorflow as mtf | |
| import tensorflow.compat.v1 as tf | |
| import random | |
| BASE_FNS = {'gelu': mtf.gelu, | |
| 'relu': mtf.relu, | |
| 'sigmoid': mtf.sigmoid, | |
| 'tanh': mtf.tanh, | |
| 'selu': mtf.selu, | |
| 'elu': mtf.elu, | |
| 'abs': mtf.abs, | |
| 'sin': mtf.sin, | |
| 'cos': mtf.cos, | |
| 'sign': mtf.sign, | |
| 'silu': mtf.swish, | |
| 'softplus': mtf.softplus | |
| } | |
| def _arcsinh(x): | |
| return mtf.log(x + mtf.sqrt(1 + x ** 2)) | |
| def _var(x, init): | |
| return mtf.get_variable(x.mesh, f"activation-{random.randint(0, 2 ** 32):x}", [], | |
| initializer=tf.constant_initializer(init), dtype=x.dtype) | |
| def _pos_var(x, val): | |
| return mtf.softplus(_var(x, 0)) + val | |
| def _rrelu(x): | |
| negative_scale = random.random() | |
| return (negative_scale * mtf.abs(x) + x) / (1 + negative_scale) | |
| def _elish(x): | |
| cond = mtf.cast(mtf.greater(x, 0), x.dtype) | |
| exp = mtf.exp(x) | |
| return cond * x / (1 + exp) + (1 - cond) * (exp - 1) / (1 / exp + 1) | |
| CUSTOM_FNS = {'lrelu001': lambda x: mtf.leaky_relu(x, alpha=0.01), | |
| 'lrelu020': lambda x: mtf.leaky_relu(x, alpha=0.20), | |
| 'id': lambda x: x, | |
| 'triangle_relax': lambda x: mtf.sin(x) - mtf.sin(3 * x) / 9 + mtf.sin(5 * x) / 25 - mtf.sin(7 * x) / 49, | |
| 'square_relax': lambda x: mtf.cos(x) - mtf.cos(3 * x) / 3 + mtf.cos(5 * x) / 5 - mtf.cos(7 * x) / 7, | |
| 'spike': lambda x: 1 / (1 + x ** 2), | |
| 'spike2': lambda x: mtf.exp(-x ** 2), | |
| 'tanhshrink': lambda x: x - tanh(x), | |
| 'softsign': lambda x: x / (mtf.abs(x) + 1), | |
| 'softmax': lambda x: mtf.softmax(x, x.shape[-1]), | |
| 'logsoftmax': lambda x: mtf.log_softmax(x, x.shape[-1]), | |
| 'bipolarsigmoid': lambda x: mtf.sigmoid(x) * 2 - 1, | |
| 'rrelu': _rrelu, | |
| 'elish': _elish, | |
| 'arcsinh': _arcsinh, | |
| 'aria': lambda x: x * (_var(x, 0) + _var(x, 1) / ( | |
| _pos_var(x, 0) + _var(x, 1) * mtf.exp(_var(x, -1) * x) ** (1 / _pos_var(x, 1)))), | |
| 'prelu': lambda x: mtf.leaky_relu(x, alpha=_var(x, 0.2)), | |
| 'parcsinh': lambda x: _var(x, 1) * _arcsinh(x * _pos_var(x, 1)), | |
| 'psoftplus': lambda x: _var(x, 1) * mtf.softplus(x * _var(x, 1)) + _var(x, 0), | |
| 'proottanh': lambda x: (x ** _pos_var(x, 2) + _pos_var(x, 1)) ** (1 / _pos_var(x, 3)) * mtf.tanh(x), | |
| 'maxsig': lambda x: mtf.maximum(x, mtf.sigmoid(x)), | |
| 'cosid': lambda x: mtf.cos(x) - x, | |
| 'minsin': lambda x: mtf.minimum(x, mtf.sin(x)), | |
| 'maxtanh': lambda x: mtf.maximum(x, mtf.tanh(x)), | |
| 'mish': lambda x: x * mtf.tanh(mtf.softplus(x)), | |
| 'tanhexp': lambda x: x * mtf.tanh(mtf.exp(x)), | |
| 'lisht': lambda x: x * mtf.tanh(x), | |
| 'seagull': lambda x: mtf.log(1 + x ** 2), | |
| 'snake': lambda x: x + mtf.sin(x) ** 2, | |
| 'roottanh': lambda x: (x ** 2 + 1) ** (1 / 3) * mtf.tanh(x), | |
| 'softplusmone': lambda x: mtf.softplus(x) - 1 | |
| } | |
| def get_activation_fn(params): | |
| if "activation_fn" in params: | |
| activation_fn = params["activation_fn"] | |
| else: | |
| print("Defaulting to GELU activation (see here: https://arxiv.org/abs/1606.08415)") | |
| activation_fn = "gelu" | |
| if activation_fn in BASE_FNS: | |
| return BASE_FNS[activation_fn] | |
| if activation_fn in CUSTOM_FNS: | |
| return CUSTOM_FNS[activation_fn] | |
| raise ValueError('unknown activation function "activation_fn" in config') | |