"Making Keras Cooler!": Elegant Layers and Fancy Callbacks

“让Keras更酷一些！”：精巧的层与花式的回调 苏剑林 · 2018-08-06 · Machine Learning model deep learning loss function Keras

Keras Has Accompanied Me All Along

Looking back at the two or three years since I entered the machine learning field, Keras has been by my side the whole time. If I hadn't stumbled upon this easy-to-use framework right when I first fell into this rabbit hole—one that let me quickly turn my ideas into code—I'm not sure I would have had the willpower to stick with it, given that back then it was still the era of theano, pylearn, caffe, torch, and the like, which even today still read to me like ancient scripture.

Later, to broaden my horizons, I spent some time learning tensorflow and wrote several programs in pure tensorflow, but no matter what, I could never quite let go of Keras. As my understanding of Keras deepened—especially after spending some time studying its source code—I found that Keras doesn't actually suffer from the "lack of flexibility" that people so often criticize it for. In fact, Keras's elegant encapsulation lets us easily implement all sorts of complex functionality. I've come to feel more and more that Keras is like an exquisite work of art, one that fully reflects the deep craftsmanship of its developers.

This post covers some aspects of custom models in Keras. Relatively speaking, this is more advanced Keras material, so beginners may want to skip it for now.

Custom Layers

Here we introduce how to build custom layers in Keras, along with some practical techniques, through which we can appreciate just how elegant Keras layers really are. more

Basic Definition Method

In Keras, the simplest way to define a custom layer is via the Lambda layer:

from keras.layers import *
from keras import backend as K

x_in = Input(shape=(10,))
x = Lambda(lambda x: x+2)(x_in) # 对输入加上2

Sometimes we want to distinguish between the training phase and the testing phase—for example, adding some noise to the input during training and removing it during testing. This requires using K.in_train_phase, like so:

def add_noise_in_train(x):
    x_ = x + K.random_normal(shape=K.shape(x)) # 加上标准高斯噪声
    return K.in_train_phase(x_, x)

x_in = Input(shape=(10,))
x = Lambda(add_noise_in_train)(x_in) # 训练阶段加入高斯噪声，测试阶段去掉

Of course, the Lambda layer is only suitable for cases where no new trainable parameters need to be introduced. If the functionality you want to implement requires adding new parameters to the model, then you must use a custom Layer. This isn't actually complicated either—compared to the Lambda layer, it just takes a few more lines of code, and the official documentation already explains it very clearly:

https://keras.io/layers/writing-your-own-keras-layers/

Here's the example from that page, reproduced below:

class MyLayer(Layer):

    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim # 可以自定义一些属性，方便调用
        super(MyLayer, self).__init__(**kwargs) # 必须

    def build(self, input_shape):
        # 添加可训练参数
        self.kernel = self.add_weight(name='kernel', 
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='uniform',
                                      trainable=True)

    def call(self, x):
        # 定义功能，相当于Lambda层的功能函数
        return K.dot(x, self.kernel)

    def compute_output_shape(self, input_shape):
        # 计算输出形状，如果输入和输出形状一致，那么可以省略，否则最好加上
        return (input_shape[0], self.output_dim)

Layers with Two Outputs

Almost every layer we normally encounter, including all the built-in layers in Keras, takes one or more inputs and returns a single output. So can Keras support layers with two outputs? The answer is yes, but you need to explicitly define output_shape. For example, the layer below simply splits the input in half and returns both halves simultaneously.

class SplitVector(Layer):

    def __init__(self, **kwargs):
        super(SplitVector, self).__init__(**kwargs)

    def call(self, inputs):
        # 按第二个维度对tensor进行切片，返回一个list
        in_dim = K.int_shape(inputs)[-1]
        return [inputs[:, :in_dim//2], inputs[:, in_dim//2:]]

    def compute_output_shape(self, input_shape):
        # output_shape也要是对应的list
        in_dim = input_shape[-1]
        return [(None, in_dim//2), (None, in_dim-in_dim//2)]

x1, x2 = SplitVector()(x_in) # 使用方法

Combining Layers with the Loss

Readers who have already gone through Custom Complex Loss Functions in Keras] will know that in Keras, a loss function is basically defined as a function of y_true and y_pred. But in more complicated situations, the loss isn't merely a function of the prediction and target values—it may also involve more elaborate computations combined with weights.

Here again we'll use center loss as an example, to introduce an implementation based on a custom layer.

class Dense_with_Center_loss(Layer):

    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super(Dense_with_Center_loss, self).__init__(**kwargs)

    def build(self, input_shape):
        # 添加可训练参数
        self.kernel = self.add_weight(name='kernel',
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='glorot_normal',
                                      trainable=True)
        self.bias = self.add_weight(name='bias',
                                    shape=(self.output_dim,),
                                    initializer='zeros',
                                    trainable=True)
        self.centers = self.add_weight(name='centers',
                                       shape=(self.output_dim, input_shape[1]),
                                       initializer='glorot_normal',
                                       trainable=True)

    def call(self, inputs):
        # 对于center loss来说，返回结果还是跟Dense的返回结果一致
        # 所以还是普通的矩阵乘法加上偏置
        self.inputs = inputs
        return K.dot(inputs, self.kernel) + self.bias

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

    def loss(self, y_true, y_pred, lamb=0.5):
        # 定义完整的loss
        y_true = K.cast(y_true, 'int32') # 保证y_true的dtype为int32
        crossentropy = K.sparse_categorical_crossentropy(y_true, y_pred, from_logits=True)
        centers = K.gather(self.centers, y_true[:, 0]) # 取出样本中心
        center_loss = K.sum(K.square(centers - self.inputs), axis=1) # 计算center loss
        return crossentropy + lamb * center_loss

f_size = 2

x_in = Input(shape=(784,))
f = Dense(f_size)(x_in)

dense_center = Dense_with_Center_loss(10)
output = dense_center(f)

model = Model(x_in, output)
model.compile(loss=dense_center.loss,
              optimizer='adam',
              metrics=['sparse_categorical_accuracy'])

# 这里是y_train是类别的整数id，不用转为one hot
model.fit(x_train, y_train, epochs=10)

Fancy Callbacks

Besides modifying the model itself, there's a lot we might want to do during training—for instance, computing some metric on the validation set after each epoch and saving the best model, or reducing the learning rate after a certain number of epochs, or adjusting regularization parameters, and so on. All of this can be achieved through callbacks.

Official callbacks page: https://keras.io/callbacks/]

Saving the Best Model

In Keras, the simplest way to keep the best model according to a validation metric is via the built-in ModelCheckpoint, for example:

checkpoint = ModelCheckpoint(filepath='./best_model.weights',
                             monitor='val_acc',
                             verbose=1,
                             save_best_only=True)

model.fit(x_train,
          y_train,
          epochs=10,
          validation_data=(x_test, y_test),
          callbacks=[checkpoint])

This approach is simple, but it has one obvious drawback: the metric involved is determined by the metrics passed to compile, and any custom metric in Keras must be written as a tensor operation. In other words, if the metric you care about can't be expressed as a tensor operation (such as BLEU score, for example), then you can't write it as a metric function at all, and this approach won't work.

So here comes a universal solution: write your own callback, and compute whatever you like inside it. For example:

from keras.callbacks import Callback

def evaluate(): # 评测函数
    pred = model.predict(x_test)
    return np.mean(pred.argmax(axis=1) == y_test) # 爱算啥就算啥

# 定义Callback器，计算验证集的acc，并保存最优模型
class Evaluate(Callback):

    def __init__(self):
        self.accs = []
        self.highest = 0.

    def on_epoch_end(self, epoch, logs=None):
        acc = evaluate()
        self.accs.append(acc)
        if acc >= self.highest: # 保存最优模型权重
            self.highest = acc
            model.save_weights('best_model.weights')

        # 爱运行什么就运行什么
        print 'acc: %s, highest: %s' % (acc, self.highest)

evaluator = Evaluate()
model.fit(x_train,
          y_train,
          epochs=10,
          callbacks=[evaluator])

Adjusting Hyperparameters

During training, we might also want to fine-tune hyperparameters. The most common need is adjusting the learning rate based on the epoch, which can be easily done using LearningRateScheduler, itself a type of callback.

from keras.callbacks import LearningRateScheduler

def lr_schedule(epoch):
    # 根据epoch返回不同的学习率
    if epoch < 50:
        lr = 1e-2
    elif epoch < 80:
        lr = 1e-3
    else:
        lr = 1e-4
    return lr

lr_scheduler = LearningRateScheduler(lr_schedule)

model.fit(x_train,
          y_train,
          epochs=10,
          callbacks=[evaluator, lr_scheduler])

What about other hyperparameters? For instance, the lamb parameter in the center loss mentioned earlier, or similar regularization coefficients. In this case, we need to set lamb as a Variable, and then define a custom callback to dynamically assign its value. For example, here's a loss I once defined:

def mycrossentropy(y_true, y_pred, e=0.1):
    loss1 = K.categorical_crossentropy(y_true, y_pred)
    loss2 = K.categorical_crossentropy(K.ones_like(y_pred)/nb_classes, y_pred)
    return (1-e)*loss1 + e*loss2

If we want to dynamically change the parameter e, we can rewrite it as:

e = K.variable(0.1)

def mycrossentropy(y_true, y_pred):
    loss1 = K.categorical_crossentropy(y_true, y_pred)
    loss2 = K.categorical_crossentropy(K.ones_like(y_pred)/nb_classes, y_pred)
    return (1-e)*loss1 + e*loss2

model.compile(loss=mycrossentropy,
              optimizer='adam')

class callback4e(Callback):
    def __init__(self, e):
        self.e = e
    def on_epoch_end(self, epoch, logs={}):
        if epoch >= 100: # 100个epoch之后设为0.01 
            K.set_value(self.e, 0.01)

model.fit(x_train,
          y_train,
          epochs=10,
          callbacks=[callback4e(e)])

Note that the Callback class supports six different hook functions corresponding to different stages: on_epoch_begin, on_epoch_end, on_batch_begin, on_batch_end, on_train_begin, and on_train_end. Each function fires at a different stage (easy enough to tell from the name), and they can be combined to implement fairly sophisticated behavior. Take warmup, for example: instead of training with the target learning rate right from the start, the learning rate is gradually ramped up from zero to the target value over the first few epochs. This process can be thought of as fine-tuning a better initialization for the model. Here's some reference code:

class Evaluate(Callback):
    def __init__(self):
        self.num_passed_batchs = 0
        self.warmup_epochs = 10
    def on_batch_begin(self, batch, logs=None):
        # params是模型自动传递给Callback的一些参数
        if self.params['steps'] == None:
            self.steps_per_epoch = np.ceil(1. * self.params['samples'] / self.params['batch_size'])
        else:
            self.steps_per_epoch = self.params['steps']
        if self.num_passed_batchs < self.steps_per_epoch * self.warmup_epochs:
            # 前10个epoch中，学习率线性地从零增加到0.001
            K.set_value(self.model.optimizer.lr,
                        0.001 * (self.num_passed_batchs + 1) / self.steps_per_epoch / self.warmup_epochs)
            self.num_passed_batchs += 1

The Endless Possibilities of Keras

Keras has plenty more noteworthy tricks—for example, you can use model.add_loss to flexibly add extra loss terms, nest models within models, or simply use Keras as a lightweight top-level API on top of tensorflow, and so on. I won't go through all of them here—readers with questions or interest are welcome to leave comments and discuss.

We usually think of highly encapsulated libraries like Keras as lacking flexibility, but that's actually not the case. You have to realize that Keras doesn't simply call ready-made high-level functions from tensorflow or theano—it merely wraps a set of basic operations through its backend, and then reimplements everything (all the various layers, optimizers, and so on) using its own backend abstraction! That's precisely how it manages to support switching between different backends.

Given that level of engineering, Keras's flexibility is beyond dispute. But this flexibility is hard to appreciate from the documentation and ordinary examples alone—often you need to read the source code before you can truly sense just how impeccable Keras's design really is. To me, implementing a complex model in Keras feels like both a challenge and a creative act, and when you succeed, you find yourself marveling at the work of art you've created.