Merge pull request #87 from WorldCereal/inference-optimization

Small improvement computing encodings

Author: kvantricht

presto/inference.py

@@ -218,15 +218,19 @@ def _get_encodings(self, dl: DataLoader) -> np.ndarray:
             np.ndarray: Array containing encoded features.
         """
 
-        all_encodings = []
+        encodings = np.empty(
+            [len(dl.dataset), self.model.encoder.embedding_size],  # type: ignore[arg-type]
+            dtype=np.float32,
+        )
+
+        with torch.no_grad():
 
-        for x, dw, latlons, month, variable_mask in dl:
-            x_f, dw_f, latlons_f, month_f, variable_mask_f = [
-                t.to(device) for t in (x, dw, latlons, month, variable_mask)
-            ]
+            for i, (x, dw, latlons, month, variable_mask) in enumerate(dl):
+                x_f, dw_f, latlons_f, month_f, variable_mask_f = [
+                    t.to(device) for t in (x, dw, latlons, month, variable_mask)
+                ]
 
-            with torch.no_grad():
-                encodings = (
+                encodings[i * self.batch_size : i * self.batch_size + self.batch_size, :] = (
                     self.model.encoder(
                         x_f,
                         dynamic_world=dw_f.long(),
@@ -238,9 +242,7 @@ def _get_encodings(self, dl: DataLoader) -> np.ndarray:
                     .numpy()
                 )
 
-            all_encodings.append(encodings)
-
-        return np.concatenate(all_encodings, axis=0)
+        return encodings
 
     def extract_presto_features(self, inarr: xr.DataArray, epsg: int = 4326) -> xr.DataArray: