首页 > 解决方案 > 通过 tf.GradientTape 使用预训练模型进行迁移学习无法收敛

问题描述

我想用 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 compile/fit 函数训练可以收敛

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

代码有什么问题?我应该如何修改它?谢谢

标签: pythonkerastensorflow2.0transfer-learning

解决方案


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


推荐阅读