delete useless code

autogl/module/nas/estimator/__init__.py

@@ -1,5 +1,3 @@
-import importlib
-import os
 from .base import BaseEstimator
 
 NAS_ESTIMATOR_DICT = {}
@@ -21,10 +19,8 @@ def register_nas_estimator(name):
     return register_nas_estimator_cls
 
 
-from .one_shot import OneShotEstimator
-from .train_scratch import TrainEstimator
-from .one_shot_hardware_aware import OneShotEstimator_HardwareAware
-from .train_scratch_hardware_aware import TrainEstimator_HardwareAware
+from .one_shot import OneShotEstimator, OneShotEstimator_HardwareAware
+from .train_scratch import TrainEstimator, TrainEstimator_HardwareAware
 
 def build_nas_estimator_from_name(name: str) -> BaseEstimator:
     """

autogl/module/nas/estimator/one_shot.py

@@ -1,31 +1,11 @@
-import torch.nn as nn
 import torch.nn.functional as F
 
 from . import register_nas_estimator
 from ..space import BaseSpace
 from .base import BaseEstimator
 from ..backend import *
-from ...train.evaluation import Evaluation, Acc
-
-# @register_nas_estimator("oneshot")
-# class OneShotEstimator(BaseEstimator):
-#     """
-#     One shot estimator.
-
-#     Use model directly to get estimations.
-#     """
-
-#     def infer(self, model: BaseSpace, dataset, mask="train"):
-#         device = next(model.parameters()).device
-#         dset = dataset[0].to(device)
-#         pred = model(dset)[getattr(dset, f"{mask}_mask")]
-#         y = dset.y[getattr(dset, f"{mask}_mask")]
-#         loss = getattr(F, self.loss_f)(pred, y)
-#         # acc=sum(pred.max(1)[1]==y).item()/y.size(0)
-#         probs = F.softmax(pred, dim=1).detach().cpu().numpy()
-#         y = y.cpu()
-#         metrics = [eva.evaluate(probs, y) for eva in self.evaluation]
-#         return metrics, loss
+from ...train.evaluation import Acc
+
 
 @register_nas_estimator("oneshot")
 class OneShotEstimator(BaseEstimator):
@@ -48,8 +28,28 @@ class OneShotEstimator(BaseEstimator):
         y = label[mask]
 
         loss = getattr(F, self.loss_f)(pred, y)
-        # acc=sum(pred.max(1)[1]==y).item()/y.size(0)
         probs = F.softmax(pred, dim=1).detach().cpu().numpy()
         y = y.cpu()
         metrics = [eva.evaluate(probs, y) for eva in self.evaluation]
         return metrics, loss
+
+
+@register_nas_estimator("oneshot_hardware")
+class OneShotEstimator_HardwareAware(OneShotEstimator):
+    """
+    One shot hardware-aware estimator.
+
+    Use model directly to get estimations.
+    """
+    def __init__(self, loss_f="nll_loss", evaluation=[Acc()], hardware_evaluation="parameter"):
+        super().__init__(loss_f, evaluation)
+        self.hardware_evaluation = hardware_evaluation
+
+    def infer(self, model: BaseSpace, dataset, mask="train"):
+        metrics, loss = super().infer(model, dataset, mask)
+        if isinstance(self.hardware_evaluation, str):
+            model_info = model.get_model_info()
+            metrics.append(model_info[self.hardware_evaluation]())
+        else:
+            metrics.append(self.hardware_evaluation(model))
+        return metrics, loss

autogl/module/nas/estimator/one_shot_hardware_aware.py

@@ -1,63 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-
-from . import register_nas_estimator
-from ..space import BaseSpace
-from .base import BaseEstimator
-from ..backend import *
-from ...train.evaluation import Evaluation, Acc
-
-# @register_nas_estimator("oneshot")
-# class OneShotEstimator(BaseEstimator):
-#     """
-#     One shot estimator.
-
-#     Use model directly to get estimations.
-#     """
-
-#     def infer(self, model: BaseSpace, dataset, mask="train"):
-#         device = next(model.parameters()).device
-#         dset = dataset[0].to(device)
-#         pred = model(dset)[getattr(dset, f"{mask}_mask")]
-#         y = dset.y[getattr(dset, f"{mask}_mask")]
-#         loss = getattr(F, self.loss_f)(pred, y)
-#         # acc=sum(pred.max(1)[1]==y).item()/y.size(0)
-#         probs = F.softmax(pred, dim=1).detach().cpu().numpy()
-#         y = y.cpu()
-#         metrics = [eva.evaluate(probs, y) for eva in self.evaluation]
-#         return metrics, loss
-
-@register_nas_estimator("oneshot_hardware")
-class OneShotEstimator_HardwareAware(BaseEstimator):
-    """
-    One shot estimator.
-
-    Use model directly to get estimations.
-    """
-    def __init__(self, loss_f="nll_loss", evaluation=[Acc()], hardware_evaluation="parameter"):
-        super().__init__(loss_f, evaluation)
-        self.evaluation = evaluation
-        self.hardware_evaluation=hardware_evaluation
-
-    def infer(self, model: BaseSpace, dataset, mask="train"):
-        device = next(model.parameters()).device
-        dset = dataset[0].to(device)
-        mask=bk_mask(dset,mask)
-
-        pred = model(dset)[mask]
-        label=bk_label(dset)
-        y = label[mask]
-
-        loss = getattr(F, self.loss_f)(pred, y)
-        # acc=sum(pred.max(1)[1]==y).item()/y.size(0)
-        probs = F.softmax(pred, dim=1).detach().cpu().numpy()
-        y = y.cpu()
-        # print(model_info)
-        # print(self.hardware_evaluation)
-        metrics = [eva.evaluate(probs, y) for eva in self.evaluation]
-        if isinstance(self.hardware_evaluation, str):
-            model_info = model.get_model_info()
-            metrics.append(model_info[self.hardware_evaluation]())
-        else:
-            metrics.append(self.hardware_evaluation(model))
-        return metrics, loss

autogl/module/nas/estimator/train_scratch.py

@@ -1,10 +1,7 @@
-import torch.nn as nn
-import torch.nn.functional as F
-
 from . import register_nas_estimator
 from ..space import BaseSpace
 from .base import BaseEstimator
-from .one_shot import OneShotEstimator
+from .one_shot import OneShotEstimator, OneShotEstimator_HardwareAware
 import torch
 
 from autogl.module.train import NodeClassificationFullTrainer, Acc
@@ -22,8 +19,6 @@ class TrainEstimator(BaseEstimator):
         self.estimator = OneShotEstimator(self.loss_f, self.evaluation)
 
     def infer(self, model: BaseSpace, dataset, mask="train"):
-        # self.trainer.model=model
-        # self.trainer.device=model.device
         boxmodel = model.wrap()
         self.trainer = NodeClassificationFullTrainer(
             model=boxmodel,
@@ -49,3 +44,14 @@ class TrainEstimator(BaseEstimator):
                 return fin, 0
             else:
                 raise e
+
+
+@register_nas_estimator("scratch_hardware")
+class TrainEstimator_HardwareAware(TrainEstimator):
+    """
+    An hardware-aware estimator which trans from scratch
+    """
+
+    def __init__(self, loss_f="nll_loss", evaluation=[Acc()], hardware_evaluation="parameter"):
+        super().__init__(loss_f, evaluation)
+        self.estimator = OneShotEstimator_HardwareAware(self.loss_f, self.evaluation, hardware_evaluation)

autogl/module/nas/estimator/train_scratch_hardware_aware.py

@@ -1,53 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-
-from . import register_nas_estimator
-from ..space import BaseSpace
-from .base import BaseEstimator
-from .one_shot_hardware_aware import OneShotEstimator_HardwareAware
-import torch
-
-from autogl.module.train import NodeClassificationFullTrainer, Acc
-
-
-@register_nas_estimator("scratch_hardware")
-class TrainEstimator_HardwareAware(BaseEstimator):
-    """
-    An estimator which trans from scratch
-    """
-
-    def __init__(self, loss_f="nll_loss", evaluation=[Acc()], hardware_evaluation="parameter"):
-        super().__init__(loss_f, evaluation)
-        self.evaluation = evaluation
-        self.hardware_evaluation = hardware_evaluation
-        self.estimator = OneShotEstimator_HardwareAware(self.loss_f, self.evaluation, self.hardware_evaluation)
-
-
-    def infer(self, model: BaseSpace, dataset, mask="train"):
-        # self.trainer.model=model
-        # self.trainer.device=model.device
-        boxmodel = model.wrap()
-        self.trainer = NodeClassificationFullTrainer(
-            model=boxmodel,
-            optimizer=torch.optim.Adam,
-            lr=0.005,
-            max_epoch=300,
-            early_stopping_round=30,
-            weight_decay=5e-4,
-            device="auto",
-            init=False,
-            feval=self.evaluation,
-            loss=self.loss_f,
-            lr_scheduler_type=None,
-        )
-        try:
-            self.trainer.train(dataset)
-            with torch.no_grad():
-                return self.estimator.infer(boxmodel.model, dataset, mask)
-        except RuntimeError as e:
-            if "cuda" in str(e) or "CUDA" in str(e):
-                INF = 100
-                fin = [-INF if eva.is_higher_better else INF for eva in self.evaluation]
-                return fin, 0
-            else:
-                raise e