add all files to full process of using surrogate function. Read readme to use.

Author: Demoren1

code/dependencies/train_config.py

@@ -42,6 +42,17 @@ class TrainConfig:
 
     best_models_save_path: str
 
+    n_epochs_final: int
+    lr_final: float
+    batch_size_final: int
+    dataset_name: str
+    final_dataset_path: str
+    output_path: str
+    width: int
+    num_cells: int
+    num_workers: int
+    n_ece_bins: int
+
     seed: Optional[int] = None
 
     # Internal fields

code/inference_surrogate.py

@@ -19,6 +19,7 @@
 import gc
 from torch.utils.data import DataLoader
 from collections import deque
+import shutil
 
 # Custom imports
 import sys
@@ -98,7 +99,7 @@ def architecture_search(self):
                 dataset,
                 batch_size=self.config.batch_size_inference,
                 shuffle=False,
-                num_workers=4,
+                num_workers=self.config.num_workers,
                 collate_fn=collate_graphs,
             )
 
@@ -259,6 +260,7 @@ def select_central_models_by_clusters(self):
             plt.show()
 
     def save_models(self):
+        shutil.rmtree(self.config.best_models_save_path, ignore_errors=True)
         os.makedirs(self.config.best_models_save_path, exist_ok=True)
 
         # Сохраняем архитектуры по одной

code/old_train_models.py

@@ -0,0 +1,232 @@
+import os
+import json
+import numpy as np
+import torch
+import nni
+from torch.utils.data import SubsetRandomSampler, SequentialSampler
+from torchvision import transforms
+from torchvision.datasets import CIFAR10, CIFAR100
+from nni.nas.evaluator.pytorch import DataLoader, Classification
+
+from DartsSpace import DARTS_with_CIFAR100 as DartsSpace
+
+
+from nni.nas.space import model_context
+from tqdm import tqdm
+from IPython.display import clear_output
+from nni.nas.evaluator.pytorch import Lightning, Trainer
+
+from dependecies.data_generator import generate_arch_dicts
+from dependecies.darts_classification_module import DartsClassificationModule
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+TEST = False
+
+
+ARCHITECTURES_PATH = "/kaggle/input/second-dataset/dataset"
+MAX_EPOCHS = 60
+LEARNING_RATE = 0.025
+BATCH_SIZE = 96
+NUM_MODLES = 2000
+
+DATASET = "CIFAR100"
+
+if DATASET == "CIFAR10":
+    MEAN = [0.49139968, 0.48215827, 0.44653124]
+    STD = [0.24703233, 0.24348505, 0.26158768]
+elif DATASET == "CIFAR100":
+    MEAN = [0.5071, 0.4867, 0.4408]
+    STD = [0.2673, 0.2564, 0.2762]
+
+SEED = 228
+# random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)  # если есть GPU
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+
+def load_json_from_directory(directory_path):
+    json_data = []
+    for root, _, files in os.walk(directory_path):
+        for file in files:
+            if file.endswith('.json'):
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    try:
+                        data = json.load(f)
+                        json_data.append(data)
+                    except json.JSONDecodeError as e:
+                        print(f"Error decoding JSON from file {file_path}: {e}")
+    return json_data
+
+
+def get_data_loaders(batch_size=512):
+    """
+    Возвращает загрузчики данных для обучения и валидации.
+
+    Параметры:
+    batch_size (int): Размер батча для загрузчиков данных. По умолчанию 1024.
+
+    Возвращает:
+    tuple: Кортеж, содержащий два объекта DataLoader:
+        - search_train_loader: Загрузчик данных для обучения.
+        - search_valid_loader: Загрузчик данных для валидации.
+    """
+    transform = transforms.Compose(
+        [
+            transforms.RandomCrop(32, padding=4),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize(MEAN, STD),
+        ]
+    )
+    if DATASET == 'CIFAR10':
+        train_data = nni.trace(CIFAR10)(
+            root="./data", train=True, download=True, transform=transform
+        )
+    elif DATASET == 'CIFAR100':
+        train_data = nni.trace(CIFAR100)(
+            root="./data", train=True, download=True, transform=transform
+        )
+    num_samples = len(train_data)
+    indices = np.random.permutation(num_samples)
+    split = int(num_samples * 0.5)
+
+    search_train_loader = DataLoader(
+        train_data,
+        batch_size=batch_size,
+        num_workers=10,
+        sampler=SubsetRandomSampler(indices[:split]),
+    )
+
+    search_valid_loader = DataLoader(
+            train_data,
+        batch_size=batch_size,
+        num_workers=10,
+        sampler=SequentialSampler(indices[split:]),
+    )
+
+    return search_train_loader, search_valid_loader
+
+
+def train_model(
+    architecture, 
+    train_loader, 
+    valid_loader, 
+    max_epochs=600, 
+    learning_rate=0.025,
+    fast_dev_run=False
+):
+    with model_context(architecture):
+        if DATASET == 'CIFAR10':
+            model = DartsSpace(width=16, num_cells=10, dataset='cifar')
+        elif DATASET == 'CIFAR100':
+            model = DartsSpace(width=16, num_cells=10, dataset='cifar100')
+    
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    #if torch.cuda.device_count() > 1:
+    #    model = torch.nn.DataParallel(model)
+    model.to(device)
+
+    evaluator = Lightning(
+        DartsClassificationModule(
+            learning_rate=learning_rate,
+            weight_decay=3e-4,
+            auxiliary_loss_weight=0.4,
+            max_epochs=max_epochs
+        ),
+        trainer=Trainer(
+            gradient_clip_val=5.0,
+            max_epochs=max_epochs,
+            fast_dev_run=fast_dev_run,
+            devices=[0]
+        ),
+        train_dataloaders=train_loader#,
+        #val_dataloaders=valid_loader
+    )
+
+    evaluator.fit(model)
+    return model
+
+
+def evaluate_and_save_results(
+    model,
+    architecture,
+    model_id,  # Новый обязательный параметр для идентификации модели
+    valid_loader,
+    folder_name="results_seq_0"
+):
+    """
+    Оценивает модель на валидационном наборе данных и сохраняет результаты в JSON.
+    Аргументы:
+    model: Обученная модель
+    architecture: Архитектура модели
+    valid_loader (DataLoader): DataLoader для валидационных данных
+    model_id: Уникальный идентификатор модели
+    folder_name (str): Папка для сохранения результатов
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    os.makedirs(folder_name, exist_ok=True)
+
+    # Перенос модели на устройство и режим оценки
+    model.to(device)
+    model.eval()
+
+    valid_correct = 0
+    valid_total = 0
+    valid_preds = []
+
+    with torch.no_grad():
+        for images, labels in valid_loader:
+            # print(labels)
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            outputs = torch.softmax(outputs, dim=1)
+            valid_preds.extend(outputs.cpu().tolist())
+            _, predicted = torch.max(outputs, 1)
+            valid_correct += (predicted == labels).sum().item()
+            valid_total += labels.size(0)
+
+    valid_accuracy = valid_correct / valid_total
+
+    # Формирование результата
+    result = {
+        "architecture": architecture,
+        "valid_predictions": valid_preds,
+        "valid_accuracy": valid_accuracy,
+    }
+
+    # Генерация имени файла с использованием model_id
+    file_name = f"model_{model_id:04d}_results.json"
+    file_path = os.path.join(folder_name, file_name)
+
+    # Сохранение результатов
+    with open(file_path, "w") as f:
+        json.dump(result, f, indent=4)
+
+    print(f"Results for model_{model_id} saved to {file_path}")
+
+
+if __name__ == "__main__":
+    arch_dicts = generate_arch_dicts(NUM_MODLES)
+    arch_dicts = [tmp_arch["architecture"] for tmp_arch in arch_dicts]
+    search_train_loader, search_valid_loader = get_data_loaders(
+        batch_size=BATCH_SIZE
+    )  # Получаем загрузчики CIFAR10
+
+    for idx, architecture in enumerate(tqdm(arch_dicts)):
+        model = train_model(  # Обучаем модель
+            architecture,
+            search_train_loader,
+            search_valid_loader,
+            max_epochs=MAX_EPOCHS,
+            learning_rate=LEARNING_RATE,
+            fast_dev_run=False
+        )
+        clear_output(wait=True)
+        
+        evaluate_and_save_results(
+            model, architecture, idx, valid_loader=search_valid_loader, folder_name="results_cifar100"
+        )  # Оцениваем и сохраняем архитектуры, предсказания на тестовом наборе CIFAR10 и accuracy

code/output/ensemble_results.txt

@@ -0,0 +1,5 @@
+Ensemble Top-1 Accuracy: 9.90%
+Ensemble ECE: 0.0125
+Number of models: 2
+Model 1 Accuracy: 10.08%
+Model 2 Accuracy: 9.90%

code/output/results.txt

@@ -0,0 +1 @@
+Ensemble Top-1 Accuracy: 9.57%Ensemble ECE: 0.0078Model 1 Top-1 Accuracy: 9.83%Model 2 Top-1 Accuracy: 9.55%Model 3 Top-1 Accuracy: 9.89%Model 4 Top-1 Accuracy: 9.93%Model 5 Top-1 Accuracy: 10.04%

code/readme.md

@@ -0,0 +1,5 @@
+Чтобы запустить все этапы, выполните команду:
+
+./start_all.sh
+
+Гиперпараметры для настройки находятся в файле surrogate_hp.json

code/start_all.sh

@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+CONFIG="surrogate_hp.json"
+
+echo "=== Запуск surrogate_train.py ==="
+python3 train_surrogate.py --hyperparameters_json "$CONFIG"
+
+echo "=== Запуск inference_surrogate.py ==="
+python3 inference_surrogate.py --hyperparameters_json "$CONFIG"
+
+echo "=== Запуск train_models.py ==="
+python3 train_models.py --hyperparameters_json "$CONFIG"
+
+echo "=== Все этапы успешно завершены ==="

code/surrogate_hp.json

@@ -1,8 +1,9 @@
 {
     "seed":42,
+    "num_workers": 4,
     "dataset_path": "third_dataset/",
     "device": "cpu",
-    "developer_mode": false,
+    "developer_mode": true,
     "n_models": 1300,
 
     "upper_margin": 0.75,
@@ -13,14 +14,14 @@
     "batch_size": 8,
     "input_dim":8,
 
-    "acc_num_epochs": 40,
+    "acc_num_epochs": 10,
     "acc_lr": 1e-2,
     "acc_final_lr": 1e-5,
     "acc_dropout": 0.2,
     "acc_n_heads": 16,
     "draw_fig_acc": false,
 
-    "div_num_epochs": 25,
+    "div_num_epochs": 5,
     "div_lr": 1e-3,
     "div_final_lr": 1e-6,
     "div_dropout": 0.1,
@@ -31,11 +32,21 @@
 
     "surrogate_inference_path": "surrogate_models/",
 
-    "n_ensemble_models": 5,
+    "n_ensemble_models": 2,
     "n_models_in_pool": 128,
-    "n_models_to_generate": 5000,
-    "batch_size_inference": 8192,
-    "min_accuracy_for_pool": 0.83,
-    "plot_tsne": true,
-    "best_models_save_path": "best_models/"
+    "n_models_to_generate": 4096,
+    "batch_size_inference": 4096,
+    "min_accuracy_for_pool": 0.85,
+    "plot_tsne": false,
+    "best_models_save_path": "best_models/",
+
+    "n_epochs_final": 1,
+    "lr_final": 0.025,
+    "batch_size_final": 96,
+    "dataset_name": "CIFAR10",
+    "final_dataset_path": "final_dataset/",
+    "output_path": "output/",
+    "width": 4,
+    "num_cells": 3,
+    "n_ece_bins": 15
 }