feat: gpu support for eval

jina-ai · Oct 15, 2021 · 25edf29 · 25edf29
1 parent 02f91d8
commit 25edf29
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 9 deletions.
diff --git a/finetuner/tuner/__init__.py b/finetuner/tuner/__init__.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Dict
 
 from ..helper import AnyDNN, DocumentArrayLike, TunerReturnType, get_tuner_class
 

diff --git a/finetuner/tuner/base.py b/finetuner/tuner/base.py
@@ -172,7 +172,7 @@ def _eval(
         """Evaluate the model on given labeled data"""
         ...
 
-    def log_evaluation(self, docs, label):
+    def log_evaluation(self, docs: DocumentArrayLike, label: str):
         if self.logger.logger.isEnabledFor(logging.DEBUG):
             if label not in self._catalogs:
                 self._catalogs[label] = evaluation.extract_catalog(docs)
@@ -183,20 +183,18 @@ def log_evaluation(self, docs, label):
                 self.logger.debug(f'{label} {name}: {value}')
 
     def _get_evaluation(self, docs, catalog):
-        self._calc_embeddings(docs)
-        self._calc_embeddings(catalog)
+        self.get_embeddings(docs)
+        self.get_embeddings(catalog)
         catalog.prune()
         to_be_scored_docs = evaluation.prepare_eval_docs(docs, catalog, limit=10)
         return {
             'hits': evaluation.get_hits_at_n(to_be_scored_docs),
             'ndcg': evaluation.get_ndcg_at_n(to_be_scored_docs),
         }
 
-    def _calc_embeddings(self, docs):
-        blobs = docs.blobs
-        embeddings = self.embed_model(blobs)
-        for doc, embed in zip(docs, embeddings):
-            doc.embedding = np.array(embed)
+    @abc.abstractmethod
+    def get_embeddings(self, docs: DocumentArrayLike):
+        """Calculates and adds the embeddings for the given Documents."""
 
 
 class BaseDataset:

diff --git a/finetuner/tuner/keras/__init__.py b/finetuner/tuner/keras/__init__.py
@@ -1,5 +1,6 @@
 from typing import Dict, Optional
 
+import numpy as np
 import tensorflow as tf
 from jina.logging.profile import ProgressBar
 from tensorflow import keras
@@ -188,5 +189,11 @@ def fit(
             'metric': {'train': metrics_train, 'eval': metrics_eval},
         }
 
+    def get_embeddings(self, data: DocumentArrayLike):
+        blobs = data.blobs
+        embeddings = self.embed_model(blobs)
+        for doc, embed in zip(data, embeddings):
+            doc.embedding = np.array(embed)
+
     def save(self, *args, **kwargs):
         self.embed_model.save(*args, **kwargs)
diff --git a/finetuner/tuner/paddle/__init__.py b/finetuner/tuner/paddle/__init__.py
@@ -1,5 +1,6 @@
 from typing import Dict, Optional
 
+import numpy as np
 import paddle
 from jina.logging.profile import ProgressBar
 from paddle import nn
@@ -176,5 +177,11 @@ def fit(
             'metric': {'train': metrics_train, 'eval': metrics_eval},
         }
 
+    def get_embeddings(self, data: DocumentArrayLike):
+        blobs = data.blobs
+        embeddings = self.embed_model(paddle.Tensor(blobs))
+        for doc, embed in zip(data, embeddings):
+            doc.embedding = np.array(embed)
+
     def save(self, *args, **kwargs):
         paddle.save(self.embed_model.state_dict(), *args, **kwargs)
diff --git a/finetuner/tuner/pytorch/__init__.py b/finetuner/tuner/pytorch/__init__.py
@@ -1,5 +1,6 @@
 from typing import Dict, Optional
 
+import numpy as np
 import torch
 import torch.nn as nn
 from jina.logging.profile import ProgressBar
@@ -189,5 +190,11 @@ def fit(
             'metric': {'train': metrics_train, 'eval': metrics_eval},
         }
 
+    def get_embeddings(self, data: DocumentArrayLike):
+        blobs = data.blobs
+        embeddings = self.embed_model(torch.Tensor(blobs))
+        for doc, embed in zip(data, embeddings):
+            doc.embedding = np.array(embed)
+
     def save(self, *args, **kwargs):
         torch.save(self.embed_model.state_dict(), *args, **kwargs)