diff --git a/configs/search-opts/DARTS-NASNet-CIFAR.config b/configs/search-opts/DARTS-NASNet-CIFAR.config index b4f0756..4fb4a27 100644 --- a/configs/search-opts/DARTS-NASNet-CIFAR.config +++ b/configs/search-opts/DARTS-NASNet-CIFAR.config @@ -9,5 +9,5 @@ "momentum" : ["float", "0.9"], "nesterov" : ["bool", "1"], "criterion": ["str", "Softmax"], - "batch_size": ["int", "256"] + "batch_size": ["int", "64"] } diff --git a/lib/models/__init__.py b/lib/models/__init__.py index 456f26c..27581fa 100644 --- a/lib/models/__init__.py +++ b/lib/models/__init__.py @@ -2,6 +2,7 @@ # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 # ################################################## from os import path as osp +from typing import List, Text __all__ = ['change_key', 'get_cell_based_tiny_net', 'get_search_spaces', 'get_cifar_models', 'get_imagenet_models', \ 'obtain_model', 'obtain_search_model', 'load_net_from_checkpoint', \ @@ -42,7 +43,7 @@ def get_cell_based_tiny_net(config): # obtain the search space, i.e., a dict mapping the operation name into a python-function for this op -def get_search_spaces(xtype, name): +def get_search_spaces(xtype, name) -> List[Text]: if xtype == 'cell': from .cell_operations import SearchSpaceNames assert name in SearchSpaceNames, 'invalid name [{:}] in {:}'.format(name, SearchSpaceNames.keys()) diff --git a/lib/models/cell_searchs/search_model_darts_nasnet.py b/lib/models/cell_searchs/search_model_darts_nasnet.py index 9702f48..85c275c 100644 --- a/lib/models/cell_searchs/search_model_darts_nasnet.py +++ b/lib/models/cell_searchs/search_model_darts_nasnet.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn from copy import deepcopy +from typing import List, Text, Dict from .search_cells import NASNetSearchCell as SearchCell from .genotypes import Structure @@ -11,7 +12,7 @@ from .genotypes import Structure # The macro structure is based on NASNet class NASNetworkDARTS(nn.Module): - def __init__(self, C, N, steps, multiplier, stem_multiplier, num_classes, search_space, affine, track_running_stats): + def __init__(self, C: int, N: int, steps: int, multiplier: int, stem_multiplier: int, num_classes: int, search_space: List[Text], affine: bool, track_running_stats: bool): super(NASNetworkDARTS, self).__init__() self._C = C self._layerN = N @@ -44,31 +45,31 @@ class NASNetworkDARTS(nn.Module): self.arch_normal_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) ) self.arch_reduce_parameters = nn.Parameter( 1e-3*torch.randn(num_edge, len(search_space)) ) - def get_weights(self): + def get_weights(self) -> List[torch.nn.Parameter]: xlist = list( self.stem.parameters() ) + list( self.cells.parameters() ) xlist+= list( self.lastact.parameters() ) + list( self.global_pooling.parameters() ) xlist+= list( self.classifier.parameters() ) return xlist - def get_alphas(self): + def get_alphas(self) -> List[torch.nn.Parameter]: return [self.arch_normal_parameters, self.arch_reduce_parameters] - def show_alphas(self): + def show_alphas(self) -> Text: with torch.no_grad(): A = 'arch-normal-parameters :\n{:}'.format( nn.functional.softmax(self.arch_normal_parameters, dim=-1).cpu() ) B = 'arch-reduce-parameters :\n{:}'.format( nn.functional.softmax(self.arch_reduce_parameters, dim=-1).cpu() ) return '{:}\n{:}'.format(A, B) - def get_message(self): + def get_message(self) -> Text: string = self.extra_repr() for i, cell in enumerate(self.cells): string += '\n {:02d}/{:02d} :: {:}'.format(i, len(self.cells), cell.extra_repr()) return string - def extra_repr(self): + def extra_repr(self) -> Text: return ('{name}(C={_C}, N={_layerN}, steps={_steps}, multiplier={_multiplier}, L={_Layer})'.format(name=self.__class__.__name__, **self.__dict__)) - def genotype(self): + def genotype(self) -> Dict[Text, List]: def _parse(weights): gene = [] for i in range(self._steps): diff --git a/lib/nas_201_api/api.py b/lib/nas_201_api/api.py index 1c9e111..00ab752 100644 --- a/lib/nas_201_api/api.py +++ b/lib/nas_201_api/api.py @@ -37,9 +37,12 @@ def print_information(information, extra_info=None, show=False): if show: print('\n'.join(strings)) return strings - +""" +This is the class for API of NAS-Bench-201. +""" class NASBench201API(object): + """ The initialization function that takes the dataset file path (or a dict loaded from that path) as input. """ def __init__(self, file_path_or_dict, verbose=True): if isinstance(file_path_or_dict, str): if verbose: print('try to create the NAS-Bench-201 api from {:}'.format(file_path_or_dict)) @@ -49,6 +52,7 @@ class NASBench201API(object): file_path_or_dict = copy.deepcopy( file_path_or_dict ) else: raise ValueError('invalid type : {:} not in [str, dict]'.format(type(file_path_or_dict))) assert isinstance(file_path_or_dict, dict), 'It should be a dict instead of {:}'.format(type(file_path_or_dict)) + self.verbose = verbose # [TODO] a flag indicating whether to print more logs keys = ('meta_archs', 'arch2infos', 'evaluated_indexes') for key in keys: assert key in file_path_or_dict, 'Can not find key[{:}] in the dict'.format(key) self.meta_archs = copy.deepcopy( file_path_or_dict['meta_archs'] )