efficient weight loading

Author: jcapp4

chebai/models/base.py

@@ -261,13 +261,15 @@ def _execute(
         model_output = self(data, **data.get("model_kwargs", dict()))
         pr, tar = self._get_prediction_and_labels(data, labels, model_output)
         d = dict(data=data, labels=labels, output=model_output, preds=pr)
-        torch.save(d,"d.pt")
+        # torch.save(data["loss_kwargs"],"kloss.pt")
+        # torch.save(data["idents"],"id.pt")
         if log:
             if self.criterion is not None:
                 loss_data, loss_labels, loss_kwargs_candidates = self._process_for_loss(
                     model_output, labels, data.get("loss_kwargs", dict())
                 )
                 loss_kwargs = dict()
+                loss_kwargs['weights'] = f.create_data_weights(batchsize=len(data['idents']),dim=data['labels'].size(dim=1),weights=data["loss_kwargs"],idents=data["idents"])
                 if self.pass_loss_kwargs:
                     loss_kwargs = loss_kwargs_candidates
                 torch.save(loss_data,"loss_data.pt")

chebai/preprocessing/bin/smiles_token/tokens.txt

@@ -984,3 +984,5 @@ p
 [ClH2+]
 [BrH2+]
 [IH2+]
+[NH3]
+[OH2]

chebai/preprocessing/collate.py

@@ -108,7 +108,7 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData:
         for d in data:
             id = d["ident"]
             weight = d["weight"]
-            loss_kwargs["ident"] = weight
+            loss_kwargs[str(id)] = weight
 
 
         return XYData(

chebai/preprocessing/datasets/base.py

@@ -1184,11 +1184,11 @@ def load_processed_data(
                 data_df = self.dynamic_split_dfs[kind]
                 data = data_df.to_dict(orient="records")
                 if kind == "train":
-                    data = f.add_train_weights(data)
+                    f.init_weights()
+                    data = f.add_val_weights(data)
                 if kind == "validation":
-                    print(kind)
                     data = f.add_val_weights(data)
-                torch.save(data,"gewicht.pt")
+                # torch.save(data,"gewicht.pt")
 
                 return data
 

extras/weight_loader.py

@@ -71,16 +71,16 @@ def find_label(id:int,path="../split/splits.csv")-> str:
 
 #to do 
 # return should be a tuple of weigths matching the sequenece of the target and label tensor
-# def create_data_weights(batchsize:int,dim:int,weights:dict[str,float],idents:tuple[int,...])-> torch.tensor:
-#     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-#     weight = torch.empty(batchsize,dim,device=device)
-#     index = 0
-#     for i in idents:
-#         w = weights[str(i)]
-#         for j in range(0,dim):
-#             weight[index][j] = float(w)
-#         index = index + 1
-#     return weight
+def create_data_weights(batchsize:int,dim:int,weights:dict[str,float],idents:tuple[int,...])-> torch.tensor:
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    weight = torch.empty(batchsize,dim,device=device)
+    index = 0
+    for i in idents:
+        w = weights[str(i)]
+        for j in range(0,dim):
+            weight[index][j] = float(w)
+        index = index + 1
+    return weight
 
 def testing():
     print("hello world")
@@ -97,17 +97,21 @@ def add_val_weights(ids):
 
 
 def add_train_weights(ids):
+    it = 0
     for i in ids:
+        if it % 10000 == 0:
+            print(it)
         ident = i["ident"]
         weight = find_weight("/home/programmer/Bachelorarbeit/weights/first_it.csv",ident=ident)
         i["weight"] = weight
+        it = it +1
     return ids
 
 def check_weights(data):
     for i in data:
         print(f"({i["ident"]} , {i["weight"]}")
 
-
-#mock_init_weights()
+init_weights()
+mock_init_weights()
 # print(get_weights((233713,51990)))