From 98e702f0345dd25ec65a15c9a96444c365e0bd30 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 22 Dec 2025 06:08:07 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20TripletDataGener?= =?UTF-8?q?ator=20sampling=20to=20O(1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Pre-compute `label_to_paths` mapping in `__init__` - Replace O(N) list comprehension with O(1) dictionary lookup for positive sampling - Optimize negative sampling to use `unique_labels` (O(1) class sampling) - Reduces sampling complexity from O(N) per item to O(1) per item - Note: Negative sampling now uniformly samples a negative class first, then an image, which balances class distribution in negative samples. --- .../datagenerators/triplet_data_generator.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/deeptuner/datagenerators/triplet_data_generator.py b/deeptuner/datagenerators/triplet_data_generator.py index 18523e7..665ac71 100644 --- a/deeptuner/datagenerators/triplet_data_generator.py +++ b/deeptuner/datagenerators/triplet_data_generator.py @@ -13,6 +13,15 @@ def __init__(self, image_paths, labels, batch_size, image_size, num_classes): self.num_classes = num_classes self.label_encoder = LabelEncoder() self.encoded_labels = self.label_encoder.fit_transform(labels) + + # Pre-compute label to paths mapping for O(1) sampling + self.label_to_paths = {} + for path, label in zip(self.image_paths, self.encoded_labels): + if label not in self.label_to_paths: + self.label_to_paths[label] = [] + self.label_to_paths[label].append(path) + self.unique_labels = np.array(list(self.label_to_paths.keys())) + self.image_data_generator = ImageDataGenerator(preprocessing_function=resnet.preprocess_input) self.on_epoch_end() print(f"Initialized TripletDataGenerator with {len(self.image_paths)} images") @@ -40,12 +49,13 @@ def _generate_triplet_batch(self, batch_image_paths, batch_labels): anchor_path = batch_image_paths[i] anchor_label = batch_labels[i] - positive_path = np.random.choice( - [p for p, l in zip(self.image_paths, self.encoded_labels) if l == anchor_label] - ) - negative_path = np.random.choice( - [p for p, l in zip(self.image_paths, self.encoded_labels) if l != anchor_label] - ) + positive_path = np.random.choice(self.label_to_paths[anchor_label]) + + while True: + neg_label = np.random.choice(self.unique_labels) + if neg_label != anchor_label: + break + negative_path = np.random.choice(self.label_to_paths[neg_label]) anchor_image = load_img(anchor_path, target_size=self.image_size) positive_image = load_img(positive_path, target_size=self.image_size)