My model consist of many custom layers out of which only one layer is trainable, that is NeuralReceiver() as shown below.

class MIMOSystem(Model): # Inherits from Keras Model

    def __init__(self, training):

        super(MIMOSystem, self).__init__()
               
        self.training = training
        self.constellation = Constellation("qam", num_bits_per_symbol)
        self.mapper = Mapper(constellation=self.constellation)
        self.demapper = Demapper("app",constellation=self.constellation)
        self.binary_source = BinarySource()
        self.channel = ApplyFlatFadingChannel(add_awgn=True)
        self.neural_receiver = NeuralReceiver() # the only trainable layer
        self.encoder = encoder = LDPC5GEncoder(k, n) 
        self.decoder = LDPC5GDecoder(encoder, hard_out=True)
        self.bce = tf.keras.losses.BinaryCrossentropy(from_logits=False)
        self.acc = tf.keras.metrics.BinaryAccuracy()
   
    @tf.function
    def __call__(self, batch_size, ebno_db):

        if self.training:
            coderate = 1.0
            codewords = self.binary_source([batch_size, num_tx_ant, k])
        else:
            coderate = k/n
            bits = self.binary_source([batch_size, num_tx_ant, k])
            codewords = self.encoder(bits)
        
        x = self.mapper(codewords)
        no = ebnodb2no(ebno_db,num_bits_per_symbol,coderate)
        channel_shape = [tf.shape(x)[0], num_rx_ant, num_tx_ant]
        h = complex_normal(channel_shape)        
        y = self.channel([x, h, no])

        x_hat, no_eff = self.neural_receiver(y,h) # custom trainable layer  
    
        llr = self.demapper([x_hat, no_eff])
        
        if self.training:
            bits_hat = tf.nn.sigmoid(llr)
            loss = self.bce(codewords, bits_hat)
            acc = self.acc(codewords, bits_hat)
            return loss, acc
        else:
            bits_hat = self.decoder(llr)                       
            return bits, bits_hat 

The trainable layer NeuralReceiver() consist of few sublayers, only two mentioned to give an idea.

class NeuralReceiver(Layer):
    def __init__(self):
        
        super().__init__()
        
        self.relu_layer = relu_layer()
        self.sign_layer = sign_layer() 
       
    def __call__(self, y_, H_):

        return x_hat, no_eff

The training loop looks like this:

NUM_TRAINING_ITERATIONS = 30000

# Instantiating the MIMOSystem model for training

model = MIMOSystem(training=True)

# Minimum value of Eb/N0 [dB] for simulations
EBN0_DB_MIN = 0.0
# Maximum value of Eb/N0 [dB] for simulations
EBN0_DB_MAX = 20.0
BATCH_SIZE = 20
filepath = "training_chks/10_10_without_coding_n.tf"
cp_callback = ModelCheckpoint(filepath=filepath,
                               monitor='train_loss',
                               save_weights_only=True,
                               save_best_only=True,
                               mode='min',
                               save_freq='epoch',
                               verbose=0)
callbacks = CallbackList(cp_callback, add_history=True, model=model)
logs = {}
callbacks.on_train_begin(logs=logs)
optimizer = tf.keras.optimizers.Adam(1e-4)
train_loss_tracker = tf.keras.metrics.Mean()
for epoch in tf.range(NUM_TRAINING_ITERATIONS):
        callbacks.on_epoch_begin(epoch, logs=logs)
        ebno_db = tf.random.uniform(shape=[],minval=EBN0_DB_MIN, maxval=EBN0_DB_MAX,dtype=tf.float32)
        # Forward pass
        with tf.GradientTape() as tape:
            loss,acc = model(BATCH_SIZE, ebno_db)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        train_loss_tracker.update_state(loss)
        train_dict= {"train_loss": train_loss_tracker.result()}
        logs["train_loss"] = train_dict["train_loss"]
        if epoch % 100 == 0:
            print(f"{epoch}/{NUM_TRAINING_ITERATIONS}  Loss: {loss:.2E}  ACC: {acc:.2E}", end="\r")
        train_loss_tracker.reset_states()
        callbacks.on_epoch_end(epoch, logs=logs)
    callbacks.on_train_end(logs=logs)

model_neuralrx = MIMOSystem(training=False)
# Run one inference to build the layers and loading the weights
model_neuralrx(tf.constant(1, tf.int32), tf.constant(10.0, tf.float32))
model_neuralrx.load_weights(filepath)

Now after training the MIMOSystem model, I want to freeze the NeuralReceiver() layer and all its sublayers and add a NN decoder after the demapper in model and train it with already trained NeuralReceiver(). How do I access the NeuralReceiver() layer and freeze it? Second, After freezing the layer, it it possible to add a new layer to this model? because the model would get changed.

CodePudding user response：

#If it's the last layer then simply put the "-1" but if you don't know then write the name of the layer then
for layer in model1.layers[-1].submodules:
    layer.trainable = False

#Now append your model, after which node you wanna append your node mention that, I am appending after the last node, So I wrote -1.
x= model1.layers[-1](_input)
x = tf.keras.layers.Dense(...)(x)
...
...
...
model = tf.keras.Model(inputs, x) 

CodePudding user response：

for layer in model1.layers[-1].submodules: layer.trainable = False

#Now append your model, after which node you wanna append your node mention that, I am appending after the last node, So I wrote -1. x= model1.layers-1 x = tf.keras.layers.Dense(...)(x) ... ... ... model = tf.keras.Model(inputs, x)