Create feature_ref_loader.py

feature_ref_loader.py

@@ -0,0 +1,23 @@
+import torch
+import numpy as np
+from utils import DEVICE
+
+
+class FeatureRefLoader:
+    def __init__(self):
+        print("Feature Ref Loader init")
+
+    # TODO: The format of feature
+    def load(self, feature_ref_file_name, num_ref=5000):
+        print("Feature Ref Loader load")
+        load_ref_data = torch.load(feature_ref_file_name, map_location=DEVICE)  # cpu
+        load_ref_data = load_ref_data.to(DEVICE)
+        feature_ref = load_ref_data[np.random.permutation(load_ref_data.shape[0])][
+            :num_ref
+        ].to(DEVICE)
+        return feature_ref
+
+
+feature_two_sample_tester_ref = FeatureRefLoader().load("./feature_ref_for_test.pt")
+feature_hwt_ref = FeatureRefLoader().load("./feature_ref_HWT.pt")
+feature_mgt_ref = FeatureRefLoader().load("./feature_ref_MGT.pt")