Keras: Custom layer parameters aren't counted in model summary

[jonathanbratt]

When I construct a model using a custom layer, the parameters in the custom layer aren't counted in the model summary. For example, here's a simple model consisting of three dense layers:

library(keras)
model_custom <- keras_model_sequential()
model_custom %>% 
    layer_dense(units = 256L, input_shape = c(784L)) %>%
    custom_layer_dense(units = 256L) %>%
    layer_dense(units = 10L, activation = 'softmax')

...where custom_layer_dense is just a wrapper around a standard dense layer:

CustomDense <- R6::R6Class(
    "CustomDense",
    inherit = KerasLayer,
    lock_objects = FALSE,

    public = list(
        initialize = function(units) {
            self$units <- units
        },

        build = function(input_shape) {
            self$dense <- keras::layer_dense(units = self$units)
            super$build(input_shape)
        },

        call = function(input, mask = NULL) {
            output <- self$dense(input) 
        },

        compute_output_shape = function(input_shape) {
            input_shape[[2]] <- self$units 
            input_shape
        }
    )
)

custom_layer_dense <- function(object, 
                               units,
                               name = NULL, 
                               trainable = TRUE) {
    create_layer(CustomDense, 
                 object, 
                 args = list(units = units,
                             name = name,
                             trainable = trainable))
}

When training, this model appears to behave the same as a model using a standard dense layer. However, calling the summary method produces:

> model_custom$summary()
Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_1 (Dense)              (None, 256)               200960    
_________________________________________________________________
r_layer (RLayer)             (None, 256)               0         
_________________________________________________________________
dense_2 (Dense)              (None, 10)                2570      
=================================================================
Total params: 203,530
Trainable params: 203,530
Non-trainable params: 0
_________________________________________________________________

The middle dense layer is attributed no parameters. Also, model_custom$weights yields only the weights for the built-in layers, not the custom layer.
Am I doing something wrong in my layer construction, or is this due to limitations in theRLayer class?

[dfalbel]

This looks like a bug to me. Note that the following works in python:

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):

  def __init__(self, units, **kwargs):
    self.units = units
    super(CustomLayer, self).__init__(**kwargs)

  def build (self, input_shape):
    self.dense = tf.keras.layers.Dense(units = self.units)
    super(CustomLayer, self).build(input_shape)

  def call(self, x):
    return self.dense(x)

  def compute_output_shape(self, input_shape):
    (input_shape[0], self.units)


model = tf.keras.Sequential([
  CustomLayer(units = 10, input_shape = (5,)),
  tf.keras.layers.Dense(units = 5)
])

model.summary()

[dfalbel]

I suggest the following workaround for now:

CustomDense <- reticulate::PyClass(
  "CustomDense",
  inherit = tensorflow::tf$keras$layers$Layer,
  list(

    `__init__` = function(self, units, ...) {
      self$units <- units
      super()$`__init__`()
    },

    build = function(self, input_shape) {
      self$dense <- layer_dense(units = reticulate::py_to_r(self$units))
      super()$build(input_shape)
    },

    call = function(self, x, ...) {
      self$dense(x)
    },

    compute_output_shape = function(input_shape) {
      list(input_shape[[1]], self$units)
    }

  )
)

This required the dev version of reticulate.

[jonathanbratt]

Excellent, that works.
Thank you very much!

[jonathanbratt]

To help me prioritize my work: do you have an estimate for how long it might take to get a bug fix into the dev version?

[dfalbel]

@jonathanbratt I am considering making the PyClass approach the recommended in Keras. Still need to do more testing, but it's theoretically much more general than the actual solution (eg. you can subclass a Dense layer if you want or, a Callback with the same interface.) What do you think?

[jonathanbratt]

I could live with that. The main downside that I see is the need to reticulate::py_to_r every time we want to pass a class variable to an R function.

Also, I had just learned how to document R6 classes with roxygen. :) How would you suggest handling custom layers inside packages, if you go the PyClass route?

[dfalbel]

The main downside that I see is the need to reticulate::py_to_r every time we want to pass a class variable to an R function.

I'm pretty sure this is fixable in reticulate side. I'll take a look next week.

The workflow would be similar to the current one. Build the private class and export a function that wraps it. For example:

CustomDense <- reticulate::PyClass(
  "CustomDense",
  inherit = tensorflow::tf$keras$layers$Layer,
  list(

    `__init__` = function(self, units, ...) {
      self$units <- units
      super()$`__init__`()
    },

    build = function(self, input_shape) {
      self$dense <- layer_dense(units = reticulate::py_to_r(self$units))
      super()$build(input_shape)
    },

    call = function(self, x, ...) {
      self$dense(x)
    },

    compute_output_shape = function(input_shape) {
      list(input_shape[[1]], self$units)
    }

  )
)

#' @export
layer_custom_dense(object, units, name = NULL, trainable = NULL, weights = NULL) {
  create_layer(CustomDense, object, list(
   units = units, 
   name = name,
   trainable = trainable,
   weights = weights
  ))
}

[jonathanbratt]

In case it would be helpful to you to see examples of actual use cases, I have several R6 custom layers in an RBERT dev branch here:
https://github.com/jonathanbratt/RBERT/tree/tf2-dev/R
(See files starting with 'tf2-layer.')

[willi3by]

@dfalbel @jonathanbratt is there now a method implemented to solve this issue with R6 custom layers? I apologize if I am missing something.

Thanks!

[jonathanbratt]

Hi @willi3by.
This issue is not solved in R6 custom layers, but in the dev version of keras there is a new function (keras::Layer) for creating custom layers which does solve this problem. I think there is also a new vignette with examples of how to use the new Layer function.

[willi3by]

Thank you for the information!

[anugrahasinha]

Although this if for R, however, python also has the issue quite similar to this.
Ref : https://github.com/tensorflow/tensorflow/issues/28748