python - 使用 tf.GradientTape 预训练模型进行迁移学习无法收敛

Question

我想使用 keras 的预训练模型进行迁移学习

import tensorflow as tf
from tensorflow import keras

base_model = keras.applications.MobileNetV2(input_shape=(96, 96, 3), include_top=False, pooling='avg')
x = base_model.outputs[0]
outputs = layers.Dense(10, activation=tf.nn.softmax)(x)

model = keras.Model(inputs=base_model.inputs, outputs=outputs)

使用keras编译/拟合函数进行训练可以收敛

model.compile(optimizer=keras.optimizers.Adam(), loss=keras.losses.SparseCategoricalCrossentropy(), metrics=['accuracy'])

history = model.fit(train_data, epochs=1)

结果是:损失:0.4402 - 准确度:0.8548

我想用 tf.GradientTape 训练，但它无法收敛

optimizer = keras.optimizers.Adam()
train_loss = keras.metrics.Mean()
train_acc = keras.metrics.SparseCategoricalAccuracy()
def train_step(data, labels):    
    with tf.GradientTape() as gt:
        pred = model(data)
        loss = keras.losses.SparseCategoricalCrossentropy()(labels, pred)

    grads = gt.gradient(loss, model.trainable_variables)

    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    train_loss(loss)
    train_acc(labels, pred)

for xs, ys in train_data:
    train_step(xs, ys)

print('train_loss = {:.3f}, train_acc = {:.3f}'.format(train_loss.result(), train_acc.result()))

但结果是:train_loss = 7.576，train_acc = 0.101

如果我只通过设置训练最后一层

base_model.trainable = False

收敛，结果为:train_loss = 0.525，train_acc = 0.823

代码有什么问题吗？我应该如何修改呢？谢谢

Answer 1

尝试使用 RELU 作为激活函数。如果您使用 RELU 以外的激活函数，则可能会出现梯度消失问题。