modify train models for preparing dataset

Author: Demoren1

code/darts_classification_module.py

@@ -0,0 +1,76 @@
+import torch
+from nni.nas.evaluator.pytorch import ClassificationModule
+from torch.nn import DataParallel
+
+
+class DartsClassificationModule(ClassificationModule):
+    def __init__(
+        self,
+        learning_rate: float = 0.001,
+        weight_decay: float = 0.,
+        auxiliary_loss_weight: float = 0.4,
+        max_epochs: int = 600
+    ):
+        self.auxiliary_loss_weight = auxiliary_loss_weight
+        # Training length will be used in LR scheduler
+        self.max_epochs = max_epochs
+        super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False, num_classes=10)
+        
+    def configure_optimizers(self):
+        """Customized optimizer with momentum, as well as a scheduler."""
+        optimizer = torch.optim.SGD(
+            self.parameters(),
+            momentum=0.9,
+            lr=self.hparams.learning_rate,
+            weight_decay=self.hparams.weight_decay
+        )
+        # Cosine annealing scheduler with T_max equal to total epochs
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=self.max_epochs, eta_min=1e-3
+        )
+        return {
+            'optimizer': optimizer,
+            'lr_scheduler': {
+                'scheduler': scheduler,
+                'interval': 'epoch',
+                'frequency': 1
+            }
+        }
+
+    def training_step(self, batch, batch_idx):
+        """Training step, customized with auxiliary loss and flexible unpacking."""
+        x, y = batch
+        out = self(x)
+        # Handle auxiliary output if present
+        if self.auxiliary_loss_weight and isinstance(out, (tuple, list)) and len(out) == 2:
+            y_hat, y_aux = out
+            loss_main = self.criterion(y_hat, y)
+            loss_aux = self.criterion(y_aux, y)
+            self.log('train_loss_main', loss_main)
+            self.log('train_loss_aux', loss_aux)
+            loss = loss_main + self.auxiliary_loss_weight * loss_aux
+        else:
+            # single output or no auxiliary
+            y_hat = out[0] if isinstance(out, (tuple, list)) else out
+            loss = self.criterion(y_hat, y)
+        self.log('train_loss', loss, prog_bar=True)
+        for name, metric in self.metrics.items():
+            self.log('train_' + name, metric(y_hat, y), prog_bar=True)
+        return loss
+
+    def on_train_epoch_start(self):
+        # Handle DataParallel wrapper when adjusting drop path
+        model = self.trainer.model
+        if isinstance(model, DataParallel):
+            target_model = model.module
+        else:
+            target_model = model
+
+        # Set drop path probability before every epoch, scaled by epoch ratio
+        if hasattr(target_model, 'set_drop_path_prob') and hasattr(target_model, 'drop_path_prob'):
+            drop_prob = target_model.drop_path_prob * self.current_epoch / self.max_epochs
+            target_model.set_drop_path_prob(drop_prob)
+
+        # Logging learning rate at the beginning of every epoch
+        lr = self.trainer.optimizers[0].param_groups[0]['lr']
+        self.log('lr', lr)

code/data_generator.ipynb

@@ -36,7 +36,7 @@
     "# !pip install random --quiet --break-system-packages\n",
     "# !pip install wandb --quiet --break-system-packages\n",
     "# !pip install pytorch-lightning --quiet --break-system-packages\n",
-    "!pip install torchmetrics --quiet --break-system-packages"
+    "# !pip install torchmetrics --quiet --break-system-packages"
    ]
   },
   {
@@ -68,6 +68,8 @@
     "from nni.nas.evaluator.pytorch import ClassificationModule\n",
     "from nni.nas.evaluator.pytorch import Lightning, Trainer\n",
     "\n",
+    "from darts_classification_module import DartsClassificationModule\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
     "from IPython.display import clear_output\n",
     "\n",
@@ -206,94 +208,6 @@
     "arch_dicts = load_json_from_directory('../home/best_models_greed_cluster')"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2025-05-09T15:55:45.281299Z",
-     "iopub.status.busy": "2025-05-09T15:55:45.280801Z",
-     "iopub.status.idle": "2025-05-09T15:55:45.290522Z",
-     "shell.execute_reply": "2025-05-09T15:55:45.289841Z",
-     "shell.execute_reply.started": "2025-05-09T15:55:45.281274Z"
-    },
-    "trusted": true
-   },
-   "outputs": [],
-   "source": [
-    "class DartsClassificationModule(ClassificationModule):\n",
-    "    def __init__(\n",
-    "        self,\n",
-    "        learning_rate: float = 0.001,\n",
-    "        weight_decay: float = 0.,\n",
-    "        auxiliary_loss_weight: float = 0.4,\n",
-    "        max_epochs: int = 600\n",
-    "    ):\n",
-    "        self.auxiliary_loss_weight = auxiliary_loss_weight\n",
-    "        # Training length will be used in LR scheduler\n",
-    "        self.max_epochs = max_epochs\n",
-    "        super().__init__(learning_rate=learning_rate, weight_decay=weight_decay, export_onnx=False, num_classes=10)\n",
-    "        \n",
-    "    def configure_optimizers(self):\n",
-    "        \"\"\"Customized optimizer with momentum, as well as a scheduler.\"\"\"\n",
-    "        optimizer = torch.optim.SGD(\n",
-    "            self.parameters(),\n",
-    "            momentum=0.9,\n",
-    "            lr=self.hparams.learning_rate,\n",
-    "            weight_decay=self.hparams.weight_decay\n",
-    "        )\n",
-    "        # Cosine annealing scheduler with T_max equal to total epochs\n",
-    "        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
-    "            optimizer, T_max=self.max_epochs, eta_min=1e-3\n",
-    "        )\n",
-    "        return {\n",
-    "            'optimizer': optimizer,\n",
-    "            'lr_scheduler': {\n",
-    "                'scheduler': scheduler,\n",
-    "                'interval': 'epoch',\n",
-    "                'frequency': 1\n",
-    "            }\n",
-    "        }\n",
-    "\n",
-    "    def training_step(self, batch, batch_idx):\n",
-    "        \"\"\"Training step, customized with auxiliary loss and flexible unpacking.\"\"\"\n",
-    "        x, y = batch\n",
-    "        out = self(x)\n",
-    "        # Handle auxiliary output if present\n",
-    "        if self.auxiliary_loss_weight and isinstance(out, (tuple, list)) and len(out) == 2:\n",
-    "            y_hat, y_aux = out\n",
-    "            loss_main = self.criterion(y_hat, y)\n",
-    "            loss_aux = self.criterion(y_aux, y)\n",
-    "            self.log('train_loss_main', loss_main)\n",
-    "            self.log('train_loss_aux', loss_aux)\n",
-    "            loss = loss_main + self.auxiliary_loss_weight * loss_aux\n",
-    "        else:\n",
-    "            # single output or no auxiliary\n",
-    "            y_hat = out[0] if isinstance(out, (tuple, list)) else out\n",
-    "            loss = self.criterion(y_hat, y)\n",
-    "        self.log('train_loss', loss, prog_bar=True)\n",
-    "        for name, metric in self.metrics.items():\n",
-    "            self.log('train_' + name, metric(y_hat, y), prog_bar=True)\n",
-    "        return loss\n",
-    "\n",
-    "    def on_train_epoch_start(self):\n",
-    "        # Handle DataParallel wrapper when adjusting drop path\n",
-    "        model = self.trainer.model\n",
-    "        if isinstance(model, DataParallel):\n",
-    "            target_model = model.module\n",
-    "        else:\n",
-    "            target_model = model\n",
-    "\n",
-    "        # Set drop path probability before every epoch, scaled by epoch ratio\n",
-    "        if hasattr(target_model, 'set_drop_path_prob') and hasattr(target_model, 'drop_path_prob'):\n",
-    "            drop_prob = target_model.drop_path_prob * self.current_epoch / self.max_epochs\n",
-    "            target_model.set_drop_path_prob(drop_prob)\n",
-    "\n",
-    "        # Logging learning rate at the beginning of every epoch\n",
-    "        lr = self.trainer.optimizers[0].param_groups[0]['lr']\n",
-    "        self.log('lr', lr)\n"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 10,
@@ -351,7 +265,6 @@
     "            gradient_clip_val=5.0,\n",
     "            max_epochs=max_epochs,\n",
     "            fast_dev_run=fast_dev_run,\n",
-    "            precision='16-mixed',\n",
     "            logger=WandbLogger(experiment=wandb.run)\n",
     "        ),\n",
     "        train_dataloaders=train_loader,\n",
@@ -363,15 +276,6 @@
     "    return model"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "torch.set_float32_matmul_precision('medium')"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -624,8 +528,7 @@
     "    folder_name (str, необязательно): Имя папки для сохранения результатов. По умолчанию \"results\".\n",
     "    Исключения:\n",
     "    ValueError: Если количество моделей и архитектур не совпадает.\n",
-    "    Результаты:\n",
-    "    Для каждой модели создается файл JSON с результатами, содержащий:\n",
+    "    Результаты:    Для каждой модели создается файл JSON с результатами, содержащий:\n",
     "    - architecture: Архитектура модели.\n",
     "    - valid_predictions: Предсказания модели на валидационном наборе данных.\n",
     "    - valid_accuracy: Точность модели на валидационном наборе данных.\n",