Добавление потери веса к окончательной потере в Керасе

Question

Я написал специальную тренировку l oop, для которой я также вычисляю потери и градиенты с помощью tf.GradientTape ()

def grads_ds(model, inputs, y_true, cw):   
    with tf.GradientTape() as ds_tape:
        y_pred = model(inputs)

        y_true = to_categorical(y_true,NUM_CLASSES)

        loss = -1*K.sum(y_true*cw*K.log(y_pred))

    grads = ds_tape.gradient(loss,model_ds.trainable_variables,unconnected_gradients=tf.UnconnectedGradients.NONE)
    return grads

Теперь я не могу понять, как добавить вес

У меня есть предварительно обученная сеть, на выходе которой глобальное среднее значение объединяется, затем выравнивается (слой 1), а затем передается на слой 4096 F C (уровень 2) , за которым следует последний слой с 3 узлами (слой 3).

Итак, я должен добавить потерю веса только этой части, потому что я оставляю предварительно обученную часть замороженной.

Чтобы добавить потерю веса, я добавляю потерю, рассчитанную с помощью следующего фрагмента кода, к потере перекрестной энтропии

weight_decay_loss = 0.0
    for layer in model.layers:
        if layer.trainable == True:
            if len(layer.get_weights())>0):
                for i in range(2):
                    weight_decay_loss += l * K.sum(K.square(layer.get_weights()[i]))

Наконец, я делаю

total_loss = K.sum(loss + weight_decay_loss)

Это правильно как добавить потерю веса к общей потере в Керасе?