merged

Author: petr12375

code/dependencies/GCN.py

@@ -1,3 +1,7 @@
+import sys
+import os
+import json
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -21,6 +25,8 @@
 from torch_geometric.loader import DataLoader
 from torch_geometric.utils import dropout_edge
 
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+from Graph import Graph
 
 class SimpleGCN(nn.Module):
     def __init__(
@@ -232,39 +238,31 @@ def process_graph(graph):
         adj, features = CustomDataset.preprocess(adj, features)
         return graph.index, adj, features
 
-    def __init__(self, graphs, accuracies=None, use_tqdm=False):
-        self.indexes = []
-        self.adjs, self.features = [], []
+    def __init__(self, models_dict_path, accuracies=None, use_tqdm=False):
+        self.models_dict_path = models_dict_path
+
         self.accuracies = (
             torch.tensor(accuracies, dtype=torch.float)
             if accuracies is not None
             else None
         )
 
-        iterator = tqdm(graphs) if use_tqdm else graphs
-
-        results = Parallel(n_jobs=-1)(
-            delayed(CustomDataset.process_graph)(graph) for graph in iterator
-        )
-
-        for index, adj, features in results:
-            self.indexes.append(index)
-            self.adjs.append(adj)
-            self.features.append(features)
-
     def __getitem__(self, index):
-        adj = self.adjs[index]
-        features = self.features[index]
+        path = self.models_dict_path[index]
+        with path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        graph = Graph(data, index=index)
+        _, adj, features = self.process_graph(graph)
         edge_index, _ = dense_to_sparse(adj)
 
         data = Data(x=features, edge_index=edge_index)
-        data.index = self.indexes[index]
+        data.index = index
         if self.accuracies is not None:
             data.y = self.accuracies[index]
         return data
 
     def __len__(self):
-        return len(self.indexes)
+        return len(self.models_dict_path)
 
 
 class TripletGraphDataset(Dataset):

code/dependencies/train_config.py

@@ -39,7 +39,8 @@ class TrainConfig:
     batch_size_inference: int
     min_accuracy_for_pool: float
     plot_tsne: bool
-
+    
+    tmp_archs_path: str
     best_models_save_path: str
 
     prepared_dataset_path: str
@@ -62,8 +63,7 @@ class TrainConfig:
     # Internal fields
     model_accuracy: Optional[Any] = field(default=None, init=False)
     model_diversity: Optional[Any] = field(default=None, init=False)
-    models_dict: list = field(default_factory=list, init=False)
-    graphs: Optional[Any] = field(default=None, init=False)
+    models_dict_path: list = field(default_factory=list, init=False)
     diversity_matrix: Optional[np.ndarray] = field(default=None, init=False)
     discrete_diversity_matrix: Optional[np.ndarray] = field(default=None, init=False)
     base_train_dataset: Optional[Any] = field(default=None, init=False)

code/inference_surrogate.py

@@ -4,6 +4,8 @@
 import os
 import json
 import argparse
+import shutil
+from pathlib import Path
 
 from torch_geometric.utils import dense_to_sparse
 import matplotlib.pyplot as plt
@@ -81,6 +83,9 @@ def _find_dist_to_best(self, best_embs, emb):
         return dists.min().item()
 
     def architecture_search(self):
+        tmp_dir = Path(self.config.tmp_archs_path)
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+        tmp_dir.mkdir(parents=True, exist_ok=True)
 
         while len(self.config.best_models) < self.config.n_ensemble_models:
             print(
@@ -90,11 +95,18 @@ def architecture_search(self):
             # 1) Сгенерировать архитектуры
             arch_dicts = generate_arch_dicts(
                 self.config.n_models_to_generate, use_tqdm=True
-            )  # list of dicts
+            )
+
+            # 2) Сохраняем как JSON-файлы
+            arch_paths = []
+            for i, arch in enumerate(arch_dicts):
+                path = tmp_dir / f"arch_{len(os.listdir(tmp_dir)) + i}.json"
+                with path.open("w", encoding="utf-8") as f:
+                    json.dump(arch, f)
+                arch_paths.append(path)
 
-            # 2) Построить графы и датасет
-            graphs = [Graph(arch, index=i) for i, arch in enumerate(arch_dicts)]
-            dataset = CustomDataset(graphs, use_tqdm=True)
+            # 3) Датасет и DataLoader
+            dataset = CustomDataset(arch_paths, use_tqdm=True)
             loader = DataLoader(
                 dataset,
                 batch_size=self.config.batch_size_inference,
@@ -103,7 +115,7 @@ def architecture_search(self):
                 collate_fn=collate_graphs,
             )
 
-            # 3) Извлечь эмбеддинги (numpy)
+            # 4) Извлечь эмбеддинги
             with torch.no_grad():
                 emb_acc_np, _ = extract_embeddings(
                     self.config.model_accuracy, loader, device, use_tqdm=True
@@ -112,47 +124,47 @@ def architecture_search(self):
                     self.config.model_diversity, loader, device, use_tqdm=True
                 )
 
-            # 4) Фильтрация по точности
-            mask = (
-                emb_acc_np >= self.config.min_accuracy_for_pool
-            )  # numpy boolean array, shape (N,)
+            # 5) Фильтрация по точности
+            mask = emb_acc_np >= self.config.min_accuracy_for_pool
+            valid_paths = [p for p, ok in zip(arch_paths, mask) if ok]
+            valid_div_embs = emb_div_np[mask].astype(np.float16)
+            valid_accs = emb_acc_np[mask]
 
-            valid_archs = [arch for arch, ok in zip(arch_dicts, mask) if ok]
-            valid_div_embs = emb_div_np[mask].astype(np.float16)  # shape (n_valid, d)
-            valid_accs = emb_acc_np[mask]  # shape (n_valid,)
-
-            for arch, emb, acc in zip(valid_archs, valid_div_embs, valid_accs):
+            # 6) Добавляем в пул
+            for path, emb, acc in zip(valid_paths, valid_div_embs, valid_accs):
+                arch = json.loads(path.read_text(encoding="utf-8"))
                 self.config.potential_archs.append(arch)
                 self.config.potential_embeddings.append(emb)
                 self.config.potential_accuracies.append(acc)
 
-            # 6) Очистка временных переменных и сборка мусора
-            del arch_dicts, graphs, dataset, loader
+            # 7) Удаляем все файлы, которые не попали в пул
+            kept = set(valid_paths)
+            for path in arch_paths:
+                if path not in kept and path.exists():
+                    path.unlink()
+
+            # 8) Очистка мусора
+            del arch_dicts, arch_paths, dataset, loader
             del emb_acc_np, emb_div_np
             torch.cuda.empty_cache()
             gc.collect()
 
-            # 7) Если пул заполнен — выбираем наиболее разнообразные модели
+            # 9) Если пул заполнен — выбираем наиболее разнообразные модели
             while (
                 len(self.config.potential_archs) >= self.config.n_models_in_pool
                 and len(self.config.best_models) < self.config.n_ensemble_models
             ):
                 if self.config.best_embeddings:
-                    best_arr = np.stack(
-                        self.config.best_embeddings
-                    )  # shape (len(best), d)
-                    # Для каждого emb в пуле — минимальное расстояние до best_arr
+                    best_arr = np.stack(self.config.best_embeddings)
                     distances = [
                         np.min(np.linalg.norm(emb - best_arr, axis=1))
                         for emb in self.config.potential_embeddings
                     ]
                 else:
-                    # Для первой модели ничем не ограничены
                     distances = [np.inf] * len(self.config.potential_embeddings)
 
                 farthest = int(np.argmax(distances))
 
-                # Добавляем в лучшие
                 self.config.best_models.append(
                     self.config.potential_archs.pop(farthest)
                 )
@@ -164,6 +176,9 @@ def architecture_search(self):
                     f"Selected #{len(self.config.best_models)}/{self.config.n_ensemble_models}: acc={acc:.4f}, dist={distances[farthest]:.4f}"
                 )
 
+        # В конце удалим временные файлы полностью
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+
     def select_central_models_by_clusters(self):
         """
         Кластеризует potential_embeddings на K кластеров,

code/surrogate_hp.json

@@ -6,6 +6,7 @@
     "developer_mode": false,
 
     "n_models": 300,
+
     "upper_margin": 0.75,
     "lower_margin": 0.25,
     "diversity_matrix_metric": "overlap", 
@@ -16,7 +17,7 @@
 
     "acc_num_epochs": 10,
     "acc_lr": 1e-2,
-    "acc_final_lr": 1e-5,
+    "acc_final_lr": 5e-4,
     "acc_dropout": 0.2,
     "acc_n_heads": 16,
     "draw_fig_acc": false,
@@ -36,9 +37,10 @@
     "n_models_in_pool": 128,
     "n_models_to_generate": 4096,
     "batch_size_inference": 4096,
-    "min_accuracy_for_pool": 0.01,
+    "min_accuracy_for_pool": 0.5,
     "plot_tsne": false,
     "best_models_save_path": "best_models/",
+    "tmp_archs_path": "datasets/tmp_archs/",
 
     "n_epochs_final": 1,
     "lr_final": 0.025,

code/surrogate_models/model_accuracy.pth

@@ -43,29 +43,27 @@ def __init__(self, config: TrainConfig):
         self.dataset_path = Path(config.dataset_path)
 
     def load_dataset(self) -> None:
+        self.config.models_dict_path = []
         for i, file_path in enumerate(tqdm(
             self.dataset_path.rglob("*.json"), desc="Loading dataset"
         )):
-            if i >= self.config.n_models:
-                break
-            try:
-                data = json.loads(file_path.read_text(encoding="utf-8"))
-                self.config.models_dict.append(data)
-            except json.JSONDecodeError as e:
-                logger.error("Error decoding JSON from %s: %s", file_path, e)
-        if len(self.config.models_dict) < self.config.n_models:
+            self.config.models_dict_path.append(file_path)
+
+        if len(self.config.models_dict_path) < self.config.n_models:
             raise ValueError(
-                f"Only {len(self.config.models_dict)} models found, but n_models={self.config.n_models}"
+                f"Only {len(self.config.models_dict_path)} model paths found, but n_models={self.config.n_models}"
             )
 
-        self.config.models_dict = random.sample(
-            self.config.models_dict, self.config.n_models
+        self.config.models_dict_path = random.sample(
+            self.config.models_dict_path, self.config.n_models
         )
 
     def _prepare_predictions(self, num_samples: Optional[int] = None):
         preds = []
-        for data in self.config.models_dict:
-            arr = np.array(data["valid_predictions"])
+        for path in self.config.models_dict_path:
+            with path.open("r", encoding="utf-8") as f:
+                data = json.load(f)
+            arr = np.array(data["test_predictions"])
             preds.append(arr[:num_samples] if num_samples else arr)
         return preds
 
@@ -92,14 +90,18 @@ def create_discrete_diversity_matrix(self) -> None:
         D[M < lower[:, None]] = -1
         self.config.discrete_diversity_matrix = D
 
-    def convert_into_graphs(self) -> None:
-        self.config.graphs = [
-            Graph(d, index=i) for i, d in enumerate(self.config.models_dict)
-        ]
+    def get_accuracies(self):
+        accs = []
+        for path in self.config.models_dict_path:
+            with path.open("r", encoding="utf-8") as f:
+                data = json.load(f)
+            accs.append(data["test_accuracy"])
+        return accs
+
 
     def create_datasets(self) -> None:
-        accs = [d["valid_accuracy"] for d in self.config.models_dict]
-        ds = CustomDataset(self.config.graphs, accs)
+        accs = self.get_accuracies()
+        ds = CustomDataset(self.config.models_dict_path, accs)
         train_n = int(self.config.train_size * self.config.n_models)
         self.config.base_train_dataset, self.config.base_valid_dataset = random_split(
             ds, [train_n, self.config.n_models - train_n]
@@ -222,12 +224,19 @@ def save_models(self) -> None:
     config = TrainConfig(**params)
 
     trainer = SurrogateTrainer(config)
+    print("Loading models")
     trainer.load_dataset()
+    print("Getting diversuty matrix")
     trainer.get_diversity_matrix()
+    print("Creating discrete diversity matrix")
     trainer.create_discrete_diversity_matrix()
-    trainer.convert_into_graphs()
+    print("Creating datasets")
     trainer.create_datasets()
+    print("Creating dataloaders")
     trainer.create_dataloaders()
+    print("Training accuracy surrogate")
     trainer.train_accuracy_model()
+    print("Training diversity surrogate")
     trainer.train_diversity_model()
+    print("Saving models")
     trainer.save_models()