[MNT] Complete image market example

examples/example_image/example_init.py

@@ -9,8 +9,8 @@ import torch
 class Model(BaseModel):
     def __init__(self):
         dir_path = os.path.dirname(os.path.abspath(__file__))
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model = ConvModel(channel=3, n_random_features=10).to(device)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = ConvModel(channel=3, n_random_features=10).to(self.device)
         self.model.load_state_dict(torch.load(os.path.join(dir_path, "conv_model.pth")))
         self.model.eval()
 
@@ -18,6 +18,7 @@ class Model(BaseModel):
         pass
 
     def predict(self, X: np.ndarray) -> np.ndarray:
+        X = torch.Tensor(X).to(self.device)
         return self.model(X)
 
     def finetune(self, X: np.ndarray, y: np.ndarray):

examples/example_image/main.py

@@ -3,7 +3,7 @@ import torch
 import get_data
 import os
 import random
-from utils import generate_uploader, generate_user, ImageDataLoader, train
+from utils import generate_uploader, generate_user, ImageDataLoader, train, eval_prediction
 import time
 
 from learnware.market import EasyMarket, BaseUserInfo
@@ -84,10 +84,6 @@ user_senmantic = {
 }
 
 
-def eval_prediction(pred_y, target_y):
-    return 0, 0
-
-
 def prepare_data():
     if dataset == "cifar10":
         X_train, y_train, X_test, y_test = get_data.get_cifar10(data_root)
@@ -163,7 +159,8 @@ def test_search(load_market=True):
         image_market = EasyMarket()
     else:
         prepare_market()
-    logger.info("Number of items in the market:", len(image_market))
+        image_market = EasyMarket()
+    logger.info("Number of items in the market: %d" % len(image_market))
 
     for i in range(n_users):
         user_data_path = os.path.join(user_save_root, "user_%d_X.npy" % (i))
@@ -174,14 +171,15 @@ def test_search(load_market=True):
         user_info = BaseUserInfo(
             id=f"user_{i}", semantic_spec=user_senmantic, stat_info={"RKMEStatSpecification": user_stat_spec}
         )
+        logger.info("Searching Market for user: %d" % (i))
         sorted_score_list, single_learnware_list, mixture_learnware_list = image_market.search_learnware(user_info)
         l = len(sorted_score_list)
         for idx in range(min(l, 10)):
             learnware = single_learnware_list[idx]
             score = sorted_score_list[idx]
             pred_y = learnware.predict(user_data)
-            acc, loss = eval_prediction(pred_y, user_label)
-            logger.info("search rank: %d, score: %.3f, learnware_id: %s, loss: %.3f" % (idx, score, learnware.id, loss))
+            acc = eval_prediction(pred_y, user_label)
+            logger.info("search rank: %d, score: %.3f, learnware_id: %s, acc: %.3f" % (idx, score, learnware.id, acc))
 
 
 if __name__ == "__main__":

examples/example_image/utils.py

@@ -158,3 +158,13 @@ def test(test_X, test_y, model, batch_size=128):
     acc = correct / total * 100
     print("Accuracy: %.2f" % (acc))
     return acc
+
+
+def eval_prediction(pred_y, target_y):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    _, predicted = torch.max(pred_y.data, 1)
+    annos = torch.from_numpy(target_y).to(device)
+    total = annos.size(0)
+    correct = (predicted == annos).sum().item()
+    criterion = nn.CrossEntropyLoss()
+    return correct / total

learnware/market/easy.py

@@ -42,8 +42,11 @@ class EasyMarket(BaseMarket):
         if rebuild:
             logger.warning("Warning! You are trying to clear current database!")
             clear_learnware_table()
-            rmtree(C.LEARNWARE_POOL_PATH)
+            rmtree(C.learnware_pool_path)
 
+        os.makedirs(C.learnware_pool_path, exist_ok=True)
+        os.makedirs(C.learnware_zip_pool_path, exist_ok=True)
+        os.makedirs(C.learnware_folder_pool_path, exist_ok=True)
         self.learnware_list, self.learnware_zip_list, self.learnware_folder_list, self.count = load_market_from_db()
 
     def check_learnware(self, learnware: Learnware) -> bool:
@@ -232,6 +235,7 @@ class EasyMarket(BaseMarket):
         weight = torch.from_numpy(weight).reshape(-1).double().to(user_rkme.device)
 
         term1 = user_rkme.inner_prod(user_rkme)
+        # print('weight:', weight.shape, 'C:', C.shape)
         term2 = weight.T @ C
         term3 = weight.T @ K @ weight
         score = float(term1 - 2 * term2 + term3)