Train a submodel within a model

Question

Suppose I have a neural network that accepts a stream of data and processes it in some way. The source of the data is remote and thus I need to transmit the data through Internet. However there is too much data in the stream to send reasonably over Internet (about 20 MB/s, and expected Internet upload speed is about 5 MB/s). Thus I want to use an encoder to highlight specific features of the stream and reduce the amount sent over. I don't want to use an auto encoder since it seems like a waste - why encode, then decode back if I can encode, send and then process the encoded version. Thus my question - having the training stream and the desired result, how do I train both networks at the same time, with the main network's input being the output of the encoder?

Answer 1

You can build a model whit an encoder part and a "main task" part.

Then the structure and the architecture of the two parts highly depend on what you are trying to do, on the availability of data, etc.

For sure, the choice of not using an autoencoder forces you to train longer the whole model, but it is legit.

import tensorflow as tf

class MyModel(tf.keras.models.Model):
    def __init__(self, latent_dim):
        super(MyModel, self).__init__()
        self.latent_dim = latent_dim # The encoded vectors dimension

        self.encoder = tf.keras.Sequential([
        tf.keras.layers.Dense(512, activation="relu", input_dim=input_dim),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(rate=0.3),
        tf.keras.layers.Dense(256, activation="relu"),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(rate=0.3),
        tf.keras.layers.Dense(latent_dim, activation='relu'),
        ])
        
        self.main_model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation="relu", input_dim=latent_dim),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(rate=0.3),
        tf.keras.layers.Dense(output_dim, activation='softmax'),
        ])

    def call(self, x):
        encoded = self.encoder(x)
        classified = self.main_model(encoded)
        return classified

Note: All the hyperparameters here are random and you should reassign to something that makes sense for your task and fine-tune them.