TensorFlow: расчет градиентов потерь при регуляризации в зависимости от входных и выходных данных модели. ⇐ Python

[Anonymous]

Обзор
Моя модель — это кодировщик, у которого есть вход Z и выход x.
Я пытаюсь использовать total_loss, который содержит как традиционное контролируемое обучение, так и термины регуляризации. У меня есть дополнительные функции (вне сети), которые используют входные данные Z и прогнозируемые выходные данные x_pred для расчета соответствующих условий регуляризации для включения в расчет потерь.

Код: Выделить всё

# Custom training function within model class
def train_step(self, Z, x):
# Define loss object
loss_object = tf.keras.losses.MeanSquaredError()
with tf.GradientTape() as tape:
# Get encoder output
x_pred = self.encoder(Z)

# Calculate traditional supervised learning data loss
data_loss = loss_object(x, x_pred)

# Calculate regularization terms
x_hat, Z_pred = calc_reg_terms(x_pred, Z) # physics-informed function
# Calculate respective regularization losses
loss_x = loss_object(x, x_hat)
loss_z = loss_object(Z, Z_pred)

""""""

Вопрос
Каков правильный метод расчета градиента моей общей потери?
Раньше я пробовал просто складывать все члены потерь вместе, а затем брать градиент суммарных потерь.

Код: Выделить всё

### PAST METHOD ###
# Calculate total loss
total_loss = data_loss + a * loss_x + b * loss_z  # a,b -> set hyperparameters
# Get gradients
grads = tape.gradient(total_loss, self.trainable_weights)

Однако, поскольку мои loss_x и loss_z определены вне кодировщика, я боюсь, что эти потери действуют скорее как смещение в расчете total_loss потому что модель фактически работает хуже, когда эти потери добавляются к data_loss. Термин data_loss имеет четкую связь с обучаемыми весами кодера, что обеспечивает четкий расчет градиента, но то же самое нельзя легко сказать о моих терминах потерь при регуляризации.
ПРИМЕЧАНИЕ. Отслеживание каждой из этих трех потерь во время обучения показывает, что data_loss может уменьшаться с каждой проходящей эпохой обучения, но как loss_x, так и loss_z имеют тенденцию выходить на плато на ранних этапах во время обучения, отсюда и страх, что они действуют скорее как нежелательное смещение в отношении total_loss.
Как правильно затем вычислить градиенты с помощью data_loss
Как правильно затем вычислить градиенты с помощью data_loss? термины code>, loss_x и loss_z?

Подробнее здесь: https://stackoverflow.com/questions/789 ... nt-on-mode