From b05fd21709a47ae7365a3385c9f73fb7df15e19f Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Wed, 26 Apr 2023 17:47:14 -0700 Subject: [PATCH] Refactor the loss functions to initialize the VGG loss function in the init function to avoid duplicated initialization. PiperOrigin-RevId: 527424556 --- mediapipe/model_maker/python/core/utils/loss_functions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mediapipe/model_maker/python/core/utils/loss_functions.py b/mediapipe/model_maker/python/core/utils/loss_functions.py index cc17ab39..a60bd2ed 100644 --- a/mediapipe/model_maker/python/core/utils/loss_functions.py +++ b/mediapipe/model_maker/python/core/utils/loss_functions.py @@ -144,6 +144,7 @@ class ImagePerceptualQualityLoss(tf.keras.losses.Loss): """Initializes ImagePerceptualQualityLoss.""" self._loss_weight = loss_weight self._losses = {} + self._vgg_loss = VGGPerceptualLoss(self._loss_weight) self._reduction = reduction def _l1_loss( @@ -164,7 +165,7 @@ class ImagePerceptualQualityLoss(tf.keras.losses.Loss): self._loss_weight = PerceptualLossWeight() if self._loss_weight.content > 0 or self._loss_weight.style > 0: - vgg_loss = VGGPerceptualLoss(self._loss_weight)(img1, img2) + vgg_loss = self._vgg_loss(img1, img2) vgg_loss_value = tf.math.add_n(vgg_loss.values()) loss_value.append(vgg_loss_value) if self._loss_weight.l1 > 0: