在 tensorflow 2.0 中进行自定义激活 function

Question

I am trying to create a custom tanh() activation function in tensorflow to work with a particular output range that I want. I want my network to output concentration multipliers, so I figured if the output of tanh() were negative it should return a value between 0 and 1, and if it were positive to output a value between 1 and 10.

Here is what I currently have

def output_activation(x):
    # function to scale tanh activation to be 1-10 if x > 0, or 0-1 if x < 0
    return tf.cond(x >= 0, lambda: tf.math.tanh(x+0.1)*10, lambda: tf.math.tanh(x) + 1)

I believe this will work with a single value, but I want to output a vector of values, to which python throws a value error ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all() ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Tensors are immutable and, from my understanding, converting to a numpy array and back will slow down network training if I am on a GPU. What is the best way to get around this error but still keep the benefits of hardware acceleration?

Answer 1

I suggest you tf.keras.backend.switch . Here a dummy example

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
from tensorflow.keras import backend as K

def output_activation(x):
    return K.switch(x >= 0, tf.math.tanh(x+0.1)*10, tf.math.tanh(x) + 1)

X = np.random.uniform(0,1, (100,10))
y = np.random.uniform(0,1, 100)

inp = Input((10,))
x = Dense(8, activation=output_activation)(inp)
out = Dense(1)(x)

model = Model(inp, out)
model.compile('adam', 'mse')
model.fit(X,y, epochs=3)

here the running notebook: https://colab.research.google.com/drive/1T_kRNUphJt9xTjiOheTgoIGpGDZaaRAg?usp=sharing