How does one add ReLU or LeakyReLU Activation Functions in PyMC?

Hi,

I’m adjusting the example used in this link:

Here is the code (slightly modified to add biases):

with pm.Model() as BayesianNN:

    ann_input = pm.Data('ann_input', X_train, mutable=True)
    ann_output = pm.Data('ann_output', y_train, mutable=True)

    # Weights from input to hidden layer
    weights_in_1 = pm.Normal('w_in_1', 0, sigma=1,
                             shape=(X.shape[1], n_hidden),
                             initval=init_1)
    
    # Bias for first hidden layer
    bias_1 = pm.Normal('b_1', 0, sigma=1, shape=(n_hidden,), initval=np.zeros(n_hidden))


    # Weights from 1st to 2nd layer
    weights_1_2 = pm.Normal('w_1_2', 0, sigma=1,
                            shape=(n_hidden, n_hidden),
                            initval=init_2)
    
    # Bias for second hidden layer
    bias_2 = pm.Normal('b_2', 0, sigma=1, shape=(n_hidden,), initval=np.zeros(n_hidden))

    
    # Weights from hidden layer to output
    weights_2_out = pm.Normal('w_2_out', 0, sigma=1,
                              shape=(n_hidden,),
                              initval=init_out)
    # Bias for output layer
    bias_out = pm.Normal('b_out', 0, sigma=1, shape=(1,), initval=np.zeros(1))

    # Build neural-network using tanh activation function
    act_1 = pm.math.maximum(pm.math.dot(ann_input,
                                     weights_in_1) + bias_1)
    act_2 = pm.math.tanh(pm.math.dot(act_1,
                                     weights_1_2) + bias_2)
    act_out = pm.math.dot(act_2, weights_2_out) + bias_out

    sigma = pm.HalfCauchy('sigma', beta=10, initval=1)
    out = pm.Normal('out', mu=act_out, sigma=sigma, total_size=y_train.shape[0], observed=ann_output, shape=act_out.shape)

I want to swap out some of the tanh activation functions for some ones that are more suited for a regression task. I get this error when I use the pm.math.maximum to try to replicate the ReLU behaviour:

TypeError: Wrong number of inputs for maximum.make_node (got 1((,)), expected 2)

Does anyone know how to fix this?

Many thanks

I think you just need to add a 0 as the second argument of pm.math.maximum().

For example, this simple model will draw a vector of length 5 and some outputs will be 0 and some will be positive, as you’d expect with relu.

with pm.Model():
    act_1 = pm.Normal('a',shape=5)
    weights_1_2 = pm.Normal('w',shape=5)
    bias_2 = pm.Normal('b',shape=5)
    act_2 = pm.math.maximum(pm.math.dot(act_1, weights_1_2) + bias_2, 0)

pm.draw(act_2)
1 Like

Thank you very much, this has solved the issue. Do you have any idea on how to change this line to a:

pm.math.maximum(pm.math.dot(ann_input, weights_in_1) + bias_1, 0) to a Leaky ReLU activation?
I would need to mimic the behaviour f(x) = x if x > 0 or 0.01*x otherwise.

Appreciate the help!

1 Like

You can use pm.math.switch

1 Like

Or just change the 2nd argument of maximum from 0 to 0.01 * x

1 Like

Thanks, I’ve implemented it using:

    # PReLU parameters
    alpha_1 = pm.Uniform('alpha_1', lower=0, upper=0.5, shape=(n_hidden,))
    alpha_2 = pm.Uniform('alpha_2', lower=0, upper=0.5, shape=(n_hidden,))

    # Neural network architecture using PReLU
    z_1 = pm.math.dot(ann_input, weights_in_1) + bias_1
    act_1 = pm.math.switch(z_1 < 0, alpha_1 * z_1, z_1)
    
    z_2 = pm.math.dot(act_1, weights_1_2) + bias_2
    act_2 = pm.math.switch(z_2 < 0, alpha_2 * z_2, z_2)

    act_out = pm.math.dot(act_2, weights_2_out) + bias_out