How to create a custom PreprocessingLayer in TF 2.2

Question

I would like to create a custom preprocessing layer using the tf.keras.layers.experimental.preprocessing.PreprocessingLayer layer.

In this custom layer, placed after the input layer, I would like to normalize my image using tf.cast(img, tf.float32) / 255.

I tried to find some code or example showing how to create this preprocessing layer, but I couldn't find.

Please, can someone provide a full example creating and using the PreprocessingLayer layer ?

Answer 1

If you want to have a custom preprocessing layer, actually you don't need to use PreprocessingLayer. You can simply subclass Layer

Take the simplest preprocessing layer Rescaling as an example, it is under the tf.keras.layers.experimental.preprocessing.Rescaling namespace. However, if you check the actual implementation, it is just subclass Layer class Source Code Link Here but has @keras_export('keras.layers.experimental.preprocessing.Rescaling')

@keras_export('keras.layers.experimental.preprocessing.Rescaling')
class Rescaling(Layer):
  """Multiply inputs by `scale` and adds `offset`.
  For instance:
  1. To rescale an input in the `[0, 255]` range
  to be in the `[0, 1]` range, you would pass `scale=1./255`.
  2. To rescale an input in the `[0, 255]` range to be in the `[-1, 1]` range,
  you would pass `scale=1./127.5, offset=-1`.
  The rescaling is applied both during training and inference.
  Input shape:
    Arbitrary.
  Output shape:
    Same as input.
  Arguments:
    scale: Float, the scale to apply to the inputs.
    offset: Float, the offset to apply to the inputs.
    name: A string, the name of the layer.
  """

  def __init__(self, scale, offset=0., name=None, **kwargs):
    self.scale = scale
    self.offset = offset
    super(Rescaling, self).__init__(name=name, **kwargs)

  def call(self, inputs):
    dtype = self._compute_dtype
    scale = math_ops.cast(self.scale, dtype)
    offset = math_ops.cast(self.offset, dtype)
    return math_ops.cast(inputs, dtype) * scale + offset

  def compute_output_shape(self, input_shape):
    return input_shape

  def get_config(self):
    config = {
        'scale': self.scale,
        'offset': self.offset,
    }
    base_config = super(Rescaling, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

So it proves that Rescaling preprocessing is just another normal layer.

The main part is the def call(self, inputs) function. You can create whatever complicated logic to preprocess your inputs and then return.

A easier documentation about custom layer can be find here

In a nutshell, you can do the preprocessing by layer, either by Lambda which for simple operation or by subclassing Layer to achieve your goal.