From 85c3fed70adf6f129a2e50ce068ea968d94a910f Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 25 Jul 2023 12:27:03 -0700 Subject: [PATCH] Add class weights to core hyperparameters and classifier library. PiperOrigin-RevId: 550962843 --- mediapipe/model_maker/python/core/hyperparameters.py | 5 ++++- mediapipe/model_maker/python/core/tasks/classifier.py | 4 +++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/mediapipe/model_maker/python/core/hyperparameters.py b/mediapipe/model_maker/python/core/hyperparameters.py index 22471655..92e1856c 100644 --- a/mediapipe/model_maker/python/core/hyperparameters.py +++ b/mediapipe/model_maker/python/core/hyperparameters.py @@ -15,7 +15,7 @@ import dataclasses import tempfile -from typing import Optional +from typing import Mapping, Optional import tensorflow as tf @@ -36,6 +36,8 @@ class BaseHParams: steps_per_epoch: An optional integer indicate the number of training steps per epoch. If not set, the training pipeline calculates the default steps per epoch as the training dataset size divided by batch size. + class_weights: An optional mapping of indices to weights for weighting the + loss function during training. shuffle: True if the dataset is shuffled before training. export_dir: The location of the model checkpoint files. distribution_strategy: A string specifying which Distribution Strategy to @@ -57,6 +59,7 @@ class BaseHParams: batch_size: int epochs: int steps_per_epoch: Optional[int] = None + class_weights: Optional[Mapping[int, float]] = None # Dataset-related parameters shuffle: bool = False diff --git a/mediapipe/model_maker/python/core/tasks/classifier.py b/mediapipe/model_maker/python/core/tasks/classifier.py index a042c0ec..d504defb 100644 --- a/mediapipe/model_maker/python/core/tasks/classifier.py +++ b/mediapipe/model_maker/python/core/tasks/classifier.py @@ -110,7 +110,9 @@ class Classifier(custom_model.CustomModel): # dataset is exhausted even if there are epochs remaining. steps_per_epoch=None, validation_data=validation_dataset, - callbacks=self._callbacks) + callbacks=self._callbacks, + class_weight=self._hparams.class_weights, + ) def evaluate(self, data: dataset.Dataset, batch_size: int = 32) -> Any: """Evaluates the classifier with the provided evaluation dataset.