Store features as float16 instead of float64

openai/CLIP#30
woctezuma · Jan 31, 2021 · 728dc16 · 728dc16
1 parent 7045723
commit 728dc16
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/build_feature_index.py b/build_feature_index.py
@@ -57,9 +57,19 @@ def build_feature_index(
 
         if (counter % 1000) == 0:
             print("{}/{} in {:.2f} s".format(counter, num_games, time() - start))
-            np.save(feature_filename, Y_hat)
-
-    np.save(feature_filename, Y_hat)
+            np.save(
+                feature_filename,
+                np.asarray(Y_hat, dtype=np.float16),
+                allow_pickle=False,
+                fix_imports=False,
+            )
+
+    np.save(
+        feature_filename,
+        np.asarray(Y_hat, dtype=np.float16),
+        allow_pickle=False,
+        fix_imports=False,
+    )
 
     return