Automatic One Hot encoding layer for Keras

keras_onehot_encoder.py

import tensorflow.keras as k
import tensorflow.keras.backend as K


def _one_hot_layer(num_classes: int):
    """
    One hot encoding layer to save massive amounts of memory in Keras

    :param num_classes:
    :return:
    """
    return k.layers.Lambda(lambda x: K.one_hot(K.cast(x, 'int64'), num_classes))


OneHotLayer = _one_hot_layer