Wiki

Tips & Tricks

Basics

tf.keras is an apparently well integrated version of Keras, which takes advantage of most TensorFlow efficiencies and optimizations, including distributed training and model deployment, while still abstracting away coding details.

As of TensorFlow 2.0, all the use and creation of layers and models is made through tf.keras. The remaining TensorFlow framework consists of both more intricate details like certain data processing methods, as well as other tools like uncertainty estimation and model deployment.

Custom layers and models

TensorFlow allows the creation of custom layers, which can then work similarly to inbuilt ones like Linear and Conv2D. To do this, one should extend the keras.layers.Layer class, ideally defining the __init__ (initialization), build (weights definition) and call (feedforward operation). Here's an example:

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

Additionally, one can define a custom model. The advantage of doing so is to be able to edit the training (fit), evaluation (evaluation), output (predict) and saving (save and save_weights) methods. This is done by extending the tf.keras.Model class. Example:

class ResNet(tf.keras.Model):

    def __init__(self):
        super(ResNet, self).__init__()
        self.block_1 = ResNetBlock()
        self.block_2 = ResNetBlock()
        self.global_pool = layers.GlobalAveragePooling2D()
        self.classifier = Dense(num_classes)

    def call(self, inputs):
        x = self.block_1(inputs)
        x = self.block_2(x)
        x = self.global_pool(x)
        return self.classifier(x)

resnet = ResNet()
dataset = ...
resnet.fit(dataset, epochs=10)
resnet.save(filepath)

So if you're wondering, "should I use the Layer class or the Model class?", ask yourself: will I need to call fit() on it? Will I need to call save() on it? If so, go with Model. If not (either because your class is just a block in a bigger system, or because you are writing training & saving code yourself), use Layer.

To be confirmed: Custom layers might not work directly with TensorFlow's usual model constructions like the sequential model. When there are custom layers involved, we might be required to define a custom model as well, even if we don't need to change the fit method or others (we can just __init__ and call methods).

Making new Layers and Models via subclassing | TensorFlow Core

Custom training

There might be times when it's useful to define a custom training pipeline. In this case, a couple of nuances should be taken into account: