From b0fe02b5cd618d37409698d4d82ffa4a8630281c Mon Sep 17 00:00:00 2001
From: Yao Yang <yanghaiqie@outlook.com>
Date: Mon, 13 Sep 2021 13:06:52 +0800
Subject: [PATCH] add parameter for other search space and algorithm

---
 autogl/module/nas/algorithm/enas.py |  9 +++++----
 autogl/module/nas/algorithm/rl.py   | 16 +++-------------
 autogl/module/nas/utils.py          |  9 +++++++++
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/autogl/module/nas/algorithm/enas.py b/autogl/module/nas/algorithm/enas.py
index 5259e07..bbc0fdd 100644
--- a/autogl/module/nas/algorithm/enas.py
+++ b/autogl/module/nas/algorithm/enas.py
@@ -22,7 +22,7 @@ from .rl import (
     ReinforceField,
     ReinforceController,
 )
-from ....utils import get_logger
+from ....utils import get_logger, process_hardware_aware_metrics
 
 LOGGER = get_logger("ENAS")
 
@@ -80,6 +80,7 @@ class Enas(BaseNAS):
         model_wd=5e-4,
         disable_progress=True,
         device="cuda",
+        param_size_weight=0
     ):
         super().__init__(device)
         self.device = device
@@ -97,6 +98,7 @@ class Enas(BaseNAS):
         self.model_lr = model_lr
         self.model_wd = model_wd
         self.disable_progress = disable_progress
+        self.param_size_weight = param_size_weight
 
     def search(self, space: BaseSpace, dset, estimator):
         self.model = space
@@ -177,8 +179,7 @@ class Enas(BaseNAS):
         for ctrl_step in range(self.ctrl_steps_aggregate):
             self._resample()
             with torch.no_grad():
-                metric, loss = self._infer(mask="val")
-            reward = metric
+                metric, loss, reward = self._infer(mask="val")
             rewards.append(reward)
             if self.entropy_weight:
                 reward += self.entropy_weight * self.controller.sample_entropy.item()
@@ -221,4 +222,4 @@ class Enas(BaseNAS):
 
     def _infer(self, mask="train"):
         metric, loss = self.estimator.infer(self.model, self.dataset, mask=mask)
-        return metric[0], loss
+        return metric[0], loss, process_hardware_aware_metrics(metric, self.param_size_weight)
diff --git a/autogl/module/nas/algorithm/rl.py b/autogl/module/nas/algorithm/rl.py
index 14dfad6..44b82d9 100644
--- a/autogl/module/nas/algorithm/rl.py
+++ b/autogl/module/nas/algorithm/rl.py
@@ -17,7 +17,7 @@ from nni.nas.pytorch.fixed import apply_fixed_architecture
 from tqdm import tqdm
 from datetime import datetime
 import numpy as np
-from ....utils import get_logger
+from ....utils import get_logger, process_hardware_aware_metrics
 
 LOGGER = get_logger("RL_NAS")
 
@@ -287,16 +287,6 @@ class ReinforceController(nn.Module):
             sampled = sampled[0]
         return sampled
 
-
-def _process_hardware_aware_metrics(metric, weight):
-    if len(metric) == 1:
-        return metric[0]
-    elif len(metric) == 2:
-        return metric[0] - metric[1] * weight
-    else:
-        raise ValueError("only one or two metric allowed")
-
-
 @register_nas_algo("rl")
 class RL(BaseNAS):
     """
@@ -472,7 +462,7 @@ class RL(BaseNAS):
 
     def _infer(self, mask="train"):
         metric, loss = self.estimator.infer(self.arch._model, self.dataset, mask=mask)
-        return metric[0], loss, _process_hardware_aware_metrics(metric, self.param_size_weight)
+        return metric[0], loss, process_hardware_aware_metrics(metric, self.param_size_weight)
 
 
 
@@ -686,4 +676,4 @@ class GraphNasRL(BaseNAS):
 
     def _infer(self, mask="train"):
         metric, loss = self.estimator.infer(self.arch._model, self.dataset, mask=mask)
-        return metric[0], loss, _process_hardware_aware_metrics(metric, self.param_size_weight), metric[1:]
+        return metric[0], loss, process_hardware_aware_metrics(metric, self.param_size_weight), metric[1:]
diff --git a/autogl/module/nas/utils.py b/autogl/module/nas/utils.py
index 9ca8a05..6ba6fea 100644
--- a/autogl/module/nas/utils.py
+++ b/autogl/module/nas/utils.py
@@ -12,6 +12,15 @@ from nni.nas.pytorch.mutables import Mutable, InputChoice, LayerChoice
 _logger = logging.getLogger(__name__)
 
 
+def process_hardware_aware_metrics(metric, weight):
+    if len(metric) == 1:
+        return metric[0]
+    elif len(metric) == 2:
+        return metric[0] - metric[1] * weight
+    else:
+        raise ValueError("only one or two metric allowed")
+
+
 def to_device(obj, device):
     """
     Move a tensor, tuple, list, or dict onto device.