如何在Tensorflow中将预训练模型作为子类模型的一部分冻结/解冻？

import tensorflow as tf

# 定义预训练模型
pretrained_model = tf.keras.applications.MobileNetV2(include_top=False, weights='imagenet')

# 冻结预训练模型的所有层
for layer in pretrained_model.layers:
    layer.trainable = False

# 创建子类模型
class MyModel(tf.keras.Model):
    def __init__(self, pretrained_model):
        super(MyModel, self).__init__()
        self.pretrained_model = pretrained_model
        self.flatten = tf.keras.layers.Flatten()
        self.dense = tf.keras.layers.Dense(10, activation='softmax')

    def call(self, inputs):
        x = self.pretrained_model(inputs)
        x = self.flatten(x)
        x = self.dense(x)
        return x

# 创建模型实例
model = MyModel(pretrained_model)

# 解冻模型的最后几层
for layer in model.pretrained_model.layers[-10:]:
    layer.trainable = True

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(train_dataset, epochs=10)

# 注意：train_dataset是你的训练数据集，需要根据实际情况进行替换