Update: Improve performance by using batches

src/CLIP.py

@@ -13,3 +13,15 @@ class CLIPImageEncoder:
             input = self.processor(images=image_pil, return_tensors="pt")
             image_features = self.model.get_image_features(**input)
             return image_features.cpu().detach().numpy()[0]
+
+    def encode_images(self, batch):
+        images = batch["image"]
+        print(images)
+        input = self.processor(images=images, return_tensors="pt")
+        with torch.no_grad():
+            #image_features = self.model.get_image_features(**input)
+            image_features = self.model(**input).last_hidden_state[:,0].cpu()
+            print(image_features)
+            print("--------------------")
+            print(self.model.get_image_features(**input).cpu().detach().numpy()[0])
+        return {"clip_embeddings": image_features}

src/app.py

@@ -6,7 +6,7 @@ import numpy as np
 from datasets import *
 
 dataset = load_dataset("huggan/CelebA-faces")
-candidate_subset = dataset["train"].select(range(10)) # This is a small CBIR app! :D
+candidate_subset = dataset["train"].select(range(40)) # This is a small CBIR app! :D
 
 def emb_dataset(dataset):
     # This function might need to be split up, to reduce start-up time of app
@@ -19,7 +19,7 @@ def emb_dataset(dataset):
 
     ## CLIP Embeddings
     clip_model = CLIPImageEncoder()
-    dataset_with_embeddings = dataset_with_embeddings.map(lambda row: {'clip_embeddings': clip_model.encode_image(row["image"])})
+    dataset_with_embeddings = dataset_with_embeddings.map(clip_model.encode_images, batched=True, batch_size=8)
 
     # Add index
     dataset_with_embeddings.add_faiss_index(column='color_embeddings')