子类建模下如何加载预训练NeuralReceiver并冻结以训练新模型

一、保存MIMOSystem中NeuralReceiver的权重 #

训练完成MIMOSystem后，直接提取并保存neural_receiver层的权重：

# 假设已完成MIMOSystem的训练
trained_mimo = MIMOSystem(training=True)
# 训练逻辑省略...

# 仅保存可训练层NeuralReceiver的权重
trained_mimo.neural_receiver.save_weights("neural_receiver_weights.h5")

二、初始化MIMOSystem2并加载预训练权重 #

由于Keras自定义层需在首次调用后才会构建权重张量，需先触发模型一次前向传播，再加载权重：

# 初始化MIMOSystem2
mimo2 = MIMOSystem2(training=True)

# 触发模型构建（传入任意合法参数即可）
_ = mimo2(batch_size=1, ebno_db=10.0)

# 加载预训练的NeuralReceiver权重
mimo2.neural_receiver.load_weights("neural_receiver_weights.h5")

三、冻结NeuralReceiver层 #

设置trainable=False递归冻结NeuralReceiver及其所有子层，避免训练时更新权重：

# 冻结NeuralReceiver层
mimo2.neural_receiver.trainable = False

# 验证冻结状态（可选）
print(f"NeuralReceiver是否冻结: {not mimo2.neural_receiver.trainable}")
for sub_layer in mimo2.neural_receiver.layers:
    print(f"子层{sub_layer.name}是否冻结: {not sub_layer.trainable}")

四、训练MIMOSystem2 #

配置优化器，仅更新NN_decoder的可训练权重：

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
epochs = 50
batch_size = 64
target_ebno = 12.0

for epoch in range(epochs):
    with tf.GradientTape() as tape:
        loss, acc = mimo2(batch_size=batch_size, ebno_db=target_ebno)
    
    # 仅计算可训练变量（NN_decoder的权重）的梯度
    grads = tape.gradient(loss, mimo2.trainable_variables)
    optimizer.apply_gradients(zip(grads, mimo2.trainable_variables))
    
    print(f"Epoch {epoch+1} | Loss: {loss.numpy():.4f} | Accuracy: {acc.numpy():.4f}")

内容的提问来源于stack exchange，提问作者Alveena