如何冻结预训练模型中的自定义层并添加新的自定义可训练层并在子类建模中对其进行训练？

Question

我的模型由许多自定义层组成，其中只有一层是可训练的，即 NeuralReceiver()，如下所示。

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.NN_decoder = NN_decoder() # new trainable layer to be added after model is trained
        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 to be frozen 
    
        llr = self.demapper([x_hat, no_eff])

        # llr = self.NN_decoder() # new trainable layer to be added after model training
        
        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 

可训练层 NeuralReceiver() 由几个子层组成，仅提及两个以给出一个想法。

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

训练循环如下所示：

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)

现在，在训练 MIMOSystem 模型之后，我想冻结 NeuralReceiver() 层及其所有子层，并在模型中的解映射器之后添加一个 NN 解码器，并使用已经训练好的 NeuralReceiver() 对其进行训练。如何访问 NeuralReceiver() 层并冻结它？二、冻结图层后，是否可以在模型中添加新图层？因为模型会改变。

Answer 1

#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) 

Answer 2

对于 model1.layers[-1].submodules 中的层： layer.trainable = False

#Now append your model, after which node you want append your node mention that, I appending after the last node, 所以我写了-1。 x= model1.layers-1 x = tf.keras.layers.Dense(...)(x) ... ... ... 模型 = tf.keras.Model（输入，x）