Keras: Question: add trainable variable to model

[MaximilianPi]

Hi,
im trying to implement a loss function with a negative binomial likelhood. However, I dont know how to add the additional variable to the models list of trainable weights. Do I have to create a custom layer for that? Here is an example code (theta is not optimized here):

x = runif(100, -1,1)
y = rnbinom(100, mu = exp(0.5*x), size = 3)

library(keras)
library(tensorflow)
tfp = reticulate::import("tensorflow_probability")
dist = tfp$distributions

K = backend()
model = keras_model_sequential()
model %>%
  layer_dense(units = 10L, activation = "relu", use_bias = T, input_shape = c(NULL, 1L)) %>%
  layer_batch_normalization() %>%
  layer_dense(units = 1, activation = "linear")


# I want to train theta0
theta0 = k_variable(0.5, k_floatx())


self_loss = function(y_true, y_pred){
  y_hat = tf$exp(y_pred)

  theta1 = k_softplus(theta0) + eps
  theta = tf$div(k_constant(1.0, k_floatx()),theta1)

  probs = k_constant(1., k_floatx()) - tf$div(theta , (theta+y_hat))+ eps
  final = dist$NegativeBinomial(total_count = theta, probs = probs)$log_prob(y_true)
  return(k_mean(-final))
}
model$add_weight(theta0, shape())

model$layers[[1]]$trainable_weights[[3]] = theta0
model %>%
  compile(optimizer = keras::optimizer_rmsprop(0.0002), loss = self_loss)


model %>%
  fit(x = matrix(x,ncol = 1,byrow = T), y = matrix(y, byrow = T, ncol = 1), batch_size = 25L, epochs = 10)

k_get_value(theta0)

[skeydan]

Hi,

unfortunately I do not understand what you are trying to do.

We have some examples on using TFP with Keras - although I don't know if you can use something like that for your use case: See e.g. https://blogs.rstudio.com/tensorflow/posts/2019-01-08-getting-started-with-tf-probability/

Otherwise if that doesn't help/apply, and I understand approximately correctly, you might want to either create a custom layer and use add_loss and add_weight on that, or use add_loss and add_weight to add the loss resp. weight to an existing layer.

[MaximilianPi]

Otherwise if that doesn't help/apply, and I understand approximately correctly, you might want to either create a custom layer and use add_loss and add_weight on that, or use add_loss and add_weight to add the loss resp. weight to an existing layer.

I have an additional variable (here theta0) which also has to be updated/optimized but is used only in the loss function. My question is now whether I have to do a workaround by defining/using a custom layer/model or adding the variable directly to the model trainable weights.

[skeydan]

I was expecting a more detailed description of what the setup should be, but googling I found this

https://github.com/gokceneraslan/neuralnet_countmodels/blob/master/Count%20models%20with%20neuralnets.ipynb

so I assume you want to port this to R - this being the part

import os
os.environ['KERAS_BACKEND'] = 'tensorflow'

from keras.models import Model
from keras.layers import Input, Dense
from keras.callbacks import EarlyStopping, TensorBoard
from keras.optimizers import RMSprop
from keras import backend as K

K.clear_session()

inputs = Input(shape=(num_feat,))
predictions = Dense(num_out, activation=tf.exp)(inputs)
model = Model(input=inputs, output=predictions)

nb = NB(theta_init=tf.zeros([1, num_out]))

# add theta as a trainable variable to Keras model
# otherwise, keras optimizers will not update it
model.layers[-1].trainable_weights.append(nb.theta_variable)

opt = RMSprop(lr=1e-2)
model.compile(optimizer=opt,
              loss=nb.loss) #nbinom loss function

early_stop = EarlyStopping(monitor='val_loss', patience=5)
tb = TensorBoard(log_dir='./logs', histogram_freq=1)

losses = model.fit(X, Y,
                   callbacks=[early_stop, tb],
                   batch_size=256,
                   validation_split=0.2,
                   nb_epoch=5000, verbose=0)

val_hist = losses.history['val_loss']
train_hist = losses.history['loss']

plt.plot(range(len(val_hist)), val_hist, 'b.-',
         range(len(train_hist)), train_hist, 'g.-')
plt.ylabel('Loss')
_ = plt.xlabel('Steps')

I'd start by actually translating the Python code exactly and seeing how well that goes.

For the loss used

class NB(object):
    def __init__(self, theta=None, theta_init=[0.0],
                 scale_factor=1.0, scope='nbinom_loss/',
                 debug=False, **theta_kwargs):

        # for numerical stability
        self.eps = 1e-10
        self.scale_factor = scale_factor
        self.debug = debug
        self.scope = scope

        with tf.name_scope(self.scope):
            # a variable may be given by user or it can be created here
            if theta is None:
                theta = tf.Variable(theta_init, dtype=tf.float32,
                                    name='theta', **theta_kwargs)

            # keep a reference to the variable itself
            self.theta_variable = theta

            # to keep dispersion always non-negative
            self.theta = tf.nn.softplus(theta)

    def loss(self, y_true, y_pred, reduce=True):
        scale_factor = self.scale_factor
        eps = self.eps

        with tf.name_scope(self.scope):
            y_true = tf.cast(y_true, tf.float32)
            y_pred = tf.cast(y_pred, tf.float32) * scale_factor

            theta = 1.0/(self.theta+eps)

            t1 = -tf.lgamma(y_true+theta+eps) 
            t2 = tf.lgamma(theta+eps)
            t3 = tf.lgamma(y_true+1.0) 
            t4 = -(theta * (tf.log(theta+eps)))
            t5 = -(y_true * (tf.log(y_pred+eps)))
            t6 = (theta+y_true) * tf.log(theta+y_pred+eps)      

            if self.debug:
                tf.summary.histogram('t1', t1)
                tf.summary.histogram('t2', t2)
                tf.summary.histogram('t3', t3)
                tf.summary.histogram('t4', t4)
                tf.summary.histogram('t5', t5)
                tf.summary.histogram('t6', t6)

            final = t1 + t2 + t3 + t4 + t5 + t6

            if reduce:
                final = tf.reduce_mean(final)

        return final

I'd try using a closure in R, or an R6 class.

[MaximilianPi]

Thanks for the response, yeah I implemeted orignally all in tensorflow but wanted to switch to keras for simplicity.
Sry for my poor description. Basically my question was how can I add atf$Variables to the model's collection of trainable weights. I didnt realize I have to add it to a model's layer and not to the model directly.
Here the solution:

model %>% layer_dense(units = 30L, input_shape = shape(12L), activation = "relu") %>%
  layer_batch_normalization() %>%
  layer_dense(units = 30L,activation = "relu") %>%
  layer_batch_normalization() %>%
  layer_dense(units = 1L, activation = "linear")

model$layers[[length(model$layers)]]$add_weight(name = 'thetaX',
                                                shape = list(),
                                                initializer = initializer_constant(0.5),
                                                trainable = TRUE)
custom_negBin_likelihood = 
  function(y_true, y_pred){
    y_hat = tf$exp(y_pred)
    theta_0 = tf$get_default_graph()$get_tensor_by_name("thetaX:0")

    theta = tf$div(k_constant(1.0, k_floatx()),(k_softplus(theta_0) + eps))
    probs = k_constant(1., k_floatx()) - tf$div(theta , (theta+y_hat))+ eps
    final = dist$NegativeBinomial(total_count = theta, probs = probs)$log_prob(y_true)
    return(k_mean(-final))
}

 model %>%
      compile(loss = self_loss,
                 optimizer = optimizer)

That way thetaX is also updated during fitting. For model serilization, ofc, I have to remove explicity the weight/variable.

Thanks for your help!

[skeydan]

Cool, thx for posting the solution!