diff --git a/autogl/module/nas/algorithm/__init__.py b/autogl/module/nas/algorithm/__init__.py index 612102a..537f223 100644 --- a/autogl/module/nas/algorithm/__init__.py +++ b/autogl/module/nas/algorithm/__init__.py @@ -4,5 +4,6 @@ NAS algorithms from .base import BaseNAS from .darts import Darts +from .enas import Enas -__all__ = ["BaseNAS", "Darts"] +__all__ = ["BaseNAS", "Darts", "Enas"] diff --git a/autogl/module/nas/enas.py b/autogl/module/nas/algorithm/enas.py similarity index 97% rename from autogl/module/nas/enas.py rename to autogl/module/nas/algorithm/enas.py index 2cad798..3d175aa 100644 --- a/autogl/module/nas/enas.py +++ b/autogl/module/nas/algorithm/enas.py @@ -6,9 +6,9 @@ import torch import torch.nn as nn import torch.nn.functional as F -from .nas import BaseNAS -from .space import SpaceModel -from .utils import AverageMeterGroup, replace_layer_choice, replace_input_choice +from .base import BaseNAS +from ..space import BaseSpace +from ..utils import AverageMeterGroup, replace_layer_choice, replace_input_choice from nni.nas.pytorch.fixed import apply_fixed_architecture _logger = logging.getLogger(__name__) def _get_mask(sampled, total): @@ -288,10 +288,10 @@ class Enas(BaseNAS): self.ctrl_kwargs=ctrl_kwargs self.ctrl_lr=ctrl_lr - def search(self, space, dset, trainer): + def search(self, space: BaseSpace, dset, estimator): self.model = space self.dataset = dset#.to(self.device) - self.trainer = trainer + self.estimator = estimator self.model_optim = torch.optim.SGD( self.model.parameters(), lr=0.01, weight_decay=3e-4 ) @@ -313,7 +313,8 @@ class Enas(BaseNAS): self._train_controller(i) selection=self.export() - return SpaceModel(space,selection,self.device) + return space.export(selection,self.device) + def _train_model(self, epoch): self.model.train() self.controller.eval() @@ -364,5 +365,5 @@ class Enas(BaseNAS): return self.controller.resample() def _infer(self): - metric, loss = self.trainer.infer(self.model, self.dataset) + metric, loss = self.estimator.infer(self.model, self.dataset) return metric, loss diff --git a/examples/test_nas.py b/examples/test_nas.py index 22aa03c..693cd10 100644 --- a/examples/test_nas.py +++ b/examples/test_nas.py @@ -8,6 +8,7 @@ from autogl.module.model import BaseModel from autogl.module.train import NodeClassificationFullTrainer from autogl.module.nas import Darts, OneShotEstimator, SinglePathNodeClassificationSpace from autogl.module.train import Acc +from autogl.module.nas.algorithm.enas import Enas if __name__ == '__main__': dataset = build_dataset_from_name('cora') @@ -30,6 +31,7 @@ if __name__ == '__main__': feval=['acc'], loss="nll_loss", lr_scheduler_type=None,), + #nas_algorithms=[Enas()], nas_algorithms=[Darts(num_epochs=1)], nas_spaces=[SinglePathNodeClassificationSpace(hidden_dim=16, ops=[GCNConv, GCNConv])], nas_estimators=[OneShotEstimator()]