1218 lines
47 KiB
Python
1218 lines
47 KiB
Python
#####################################################
|
|
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.07 #
|
|
##############################################################################
|
|
# NATS-Bench: Benchmarking NAS algorithms for Architecture Topology and Size #
|
|
##############################################################################
|
|
"""In this file, we define NASBenchMetaAPI, ArchResults, and ResultsCount.
|
|
|
|
NASBenchMetaAPI is the abstract class for benchmark APIs.
|
|
We also define the class ArchResults, which contains all
|
|
information of a single architecture trained by one kind of hyper-parameters
|
|
on three datasets. We also define the class ResultsCount, which contains all
|
|
information of a single trial for a single architecture.
|
|
"""
|
|
import abc
|
|
import bz2
|
|
import collections
|
|
import copy
|
|
import os
|
|
import pickle
|
|
import random
|
|
import time
|
|
from typing import Any, Dict, Optional, Text, Union
|
|
import warnings
|
|
|
|
import numpy as np
|
|
|
|
|
|
_FILE_SYSTEM = 'default'
|
|
PICKLE_EXT = 'pickle.pbz2'
|
|
|
|
|
|
def time_string():
|
|
iso_time_format = '%Y-%m-%d %X'
|
|
string = '[{:}]'.format(
|
|
time.strftime(iso_time_format, time.gmtime(time.time())))
|
|
return string
|
|
|
|
|
|
def reset_file_system(lib: Text = 'default'):
|
|
global _FILE_SYSTEM
|
|
_FILE_SYSTEM = lib
|
|
|
|
|
|
def get_file_system():
|
|
return _FILE_SYSTEM
|
|
|
|
|
|
def get_torch_home():
|
|
if 'TORCH_HOME' in os.environ:
|
|
return os.environ['TORCH_HOME']
|
|
elif 'HOME' in os.environ:
|
|
return os.path.join(os.environ['HOME'], '.torch')
|
|
else:
|
|
raise ValueError('Did not find HOME in os.environ. '
|
|
'Please at least setup the path of HOME or TORCH_HOME '
|
|
'in the environment.')
|
|
|
|
|
|
def nats_is_dir(file_path):
|
|
if _FILE_SYSTEM == 'default':
|
|
return os.path.isdir(file_path)
|
|
elif _FILE_SYSTEM == 'google':
|
|
import tensorflow as tf # pylint: disable=g-import-not-at-top
|
|
return tf.io.gfile.isdir(file_path)
|
|
else:
|
|
raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))
|
|
|
|
|
|
def nats_is_file(file_path):
|
|
if _FILE_SYSTEM == 'default':
|
|
return os.path.isfile(file_path)
|
|
elif _FILE_SYSTEM == 'google':
|
|
import tensorflow as tf # pylint: disable=g-import-not-at-top
|
|
return tf.io.gfile.exists(file_path) and not tf.io.gfile.isdir(file_path)
|
|
else:
|
|
raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))
|
|
|
|
|
|
def pickle_save(obj, file_path, ext='.pbz2', protocol=4):
|
|
"""Use pickle to save data (obj) into file_path.
|
|
|
|
Args:
|
|
obj: The object to be saved into a path.
|
|
file_path: The target saving path.
|
|
ext: The extension of file name.
|
|
protocol: The pickle protocol. According to this documentation
|
|
(https://docs.python.org/3/library/pickle.html#data-stream-format),
|
|
the protocol version 4 was added in Python 3.4. It adds support for very
|
|
large objects, pickling more kinds of objects, and some data format
|
|
optimizations. It is the default protocol starting with Python 3.8.
|
|
"""
|
|
# with open(file_path, 'wb') as cfile:
|
|
if _FILE_SYSTEM == 'default':
|
|
with bz2.BZ2File(str(file_path) + ext, 'wb') as cfile:
|
|
pickle.dump(obj, cfile, protocol=protocol) # pytype: disable=wrong-arg-types
|
|
else:
|
|
raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))
|
|
|
|
|
|
def pickle_load(file_path, ext='.pbz2'):
|
|
"""Use pickle to load the file on different systems."""
|
|
# return pickle.load(open(file_path, "rb"))
|
|
if nats_is_file(str(file_path)):
|
|
xfile_path = str(file_path)
|
|
else:
|
|
xfile_path = str(file_path) + ext
|
|
if _FILE_SYSTEM == 'default':
|
|
with bz2.BZ2File(xfile_path, 'rb') as cfile:
|
|
return pickle.load(cfile) # pytype: disable=wrong-arg-types
|
|
elif _FILE_SYSTEM == 'google':
|
|
import tensorflow as tf # pylint: disable=g-import-not-at-top
|
|
file_content = tf.io.gfile.GFile(file_path, mode='rb').read()
|
|
byte_content = bz2.decompress(file_content)
|
|
return pickle.loads(byte_content)
|
|
else:
|
|
raise ValueError('Unknown file system lib: {:}'.format(_FILE_SYSTEM))
|
|
|
|
|
|
def remap_dataset_set_names(dataset, metric_on_set, verbose=False):
|
|
"""Re-map the metric_on_set to internal keys."""
|
|
if verbose:
|
|
print('Call internal function _remap_dataset_set_names with dataset={:} '
|
|
'and metric_on_set={:}'.format(dataset, metric_on_set))
|
|
if dataset == 'cifar10' and metric_on_set == 'valid':
|
|
dataset, metric_on_set = 'cifar10-valid', 'x-valid'
|
|
elif dataset == 'cifar10' and metric_on_set == 'test':
|
|
dataset, metric_on_set = 'cifar10', 'ori-test'
|
|
elif dataset == 'cifar10' and metric_on_set == 'train':
|
|
dataset, metric_on_set = 'cifar10', 'train'
|
|
elif (dataset == 'cifar100' or
|
|
dataset == 'ImageNet16-120') and metric_on_set == 'valid':
|
|
metric_on_set = 'x-valid'
|
|
elif (dataset == 'cifar100' or
|
|
dataset == 'ImageNet16-120') and metric_on_set == 'test':
|
|
metric_on_set = 'x-test'
|
|
if verbose:
|
|
print(' return dataset={:} and metric_on_set={:}'.format(
|
|
dataset, metric_on_set))
|
|
return dataset, metric_on_set
|
|
|
|
|
|
class NASBenchMetaAPI(metaclass=abc.ABCMeta):
|
|
"""The abstract class for NATS Bench API."""
|
|
|
|
@abc.abstractmethod
|
|
def __init__(self,
|
|
file_path_or_dict: Optional[Union[Text, Dict[Text, Any]]] = None,
|
|
fast_mode: bool = False,
|
|
verbose: bool = True):
|
|
"""The initialization function that takes the dataset file path (or a dict loaded from that path) as input."""
|
|
# NOTE(xuanyidong): the following attributes must be initilaized in subclass
|
|
self.meta_archs = None
|
|
self.verbose = None
|
|
self.evaluated_indexes = None
|
|
self.arch2infos_dict = None
|
|
self.filename = None
|
|
self._fast_mode = None
|
|
self._archive_dir = None
|
|
self._avaliable_hps = None
|
|
self.archstr2index = None
|
|
|
|
def __getitem__(self, index: int):
|
|
return copy.deepcopy(self.meta_archs[index])
|
|
|
|
def arch(self, index: int):
|
|
"""Return the topology structure of the `index`-th architecture."""
|
|
if self.verbose:
|
|
print('Call the arch function with index={:}'.format(index))
|
|
if index < 0 or index >= len(self.meta_archs):
|
|
raise ValueError('invalid index : {:} vs. {:}.'.format(
|
|
index, len(self.meta_archs)))
|
|
return copy.deepcopy(self.meta_archs[index])
|
|
|
|
def __len__(self):
|
|
return len(self.meta_archs)
|
|
|
|
def __repr__(self):
|
|
return ('{name}({num}/{total} architectures, fast_mode={fast_mode}, '
|
|
'file={filename})'.format(
|
|
name=self.__class__.__name__,
|
|
num=len(self.evaluated_indexes), total=len(self.meta_archs),
|
|
fast_mode=self.fast_mode, filename=self.filename))
|
|
|
|
@property
|
|
def avaliable_hps(self):
|
|
return list(copy.deepcopy(self._avaliable_hps))
|
|
|
|
@property
|
|
def used_time(self):
|
|
return self._used_time
|
|
|
|
@property
|
|
def search_space_name(self):
|
|
return self._search_space_name
|
|
|
|
@property
|
|
def fast_mode(self):
|
|
return self._fast_mode
|
|
|
|
@property
|
|
def archive_dir(self):
|
|
return self._archive_dir
|
|
|
|
@property
|
|
def full_train_epochs(self):
|
|
return self._full_train_epochs
|
|
|
|
def reset_archive_dir(self, archive_dir):
|
|
self._archive_dir = archive_dir
|
|
|
|
def reset_fast_mode(self, fast_mode):
|
|
self._fast_mode = fast_mode
|
|
|
|
def reset_time(self):
|
|
self._used_time = 0
|
|
|
|
@abc.abstractmethod
|
|
def get_more_info(self,
|
|
index,
|
|
dataset,
|
|
iepoch=None,
|
|
hp: Text = '12',
|
|
is_random: bool = True):
|
|
"""Return the metric for the `index`-th architecture."""
|
|
|
|
def simulate_train_eval(self,
|
|
arch,
|
|
dataset,
|
|
iepoch=None,
|
|
hp='12',
|
|
account_time=True):
|
|
"""This function is used to simulate training and evaluating an arch."""
|
|
index = self.query_index_by_arch(arch)
|
|
all_names = ('cifar10', 'cifar100', 'ImageNet16-120')
|
|
if dataset not in all_names:
|
|
raise ValueError('Invalid dataset name : {:} vs {:}'.format(
|
|
dataset, all_names))
|
|
if dataset == 'cifar10':
|
|
info = self.get_more_info(
|
|
index, 'cifar10-valid', iepoch=iepoch, hp=hp, is_random=True)
|
|
else:
|
|
info = self.get_more_info(
|
|
index, dataset, iepoch=iepoch, hp=hp, is_random=True)
|
|
valid_acc, time_cost = info[
|
|
'valid-accuracy'], info['train-all-time'] + info['valid-per-time']
|
|
latency = self.get_latency(index, dataset)
|
|
if account_time:
|
|
self._used_time += time_cost
|
|
return valid_acc, latency, time_cost, self._used_time
|
|
|
|
def random(self):
|
|
"""Return a random index of all architectures."""
|
|
return random.randint(0, len(self.meta_archs)-1)
|
|
|
|
def reload(self, archive_root: Text = None, index: int = None):
|
|
"""Overwrite all information of the 'index'-th architecture in search space.
|
|
|
|
Args:
|
|
archive_root: If archive_root is None, it will try to load from the
|
|
default path os.environ['TORCH_HOME'] / 'BASE_NAME'-full.
|
|
index: If index is None, overwrite all ckps.
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call clear_params with archive_root={:} and index={:}'.format(
|
|
time_string(), archive_root, index))
|
|
if archive_root is None:
|
|
archive_root = os.path.join(os.environ['TORCH_HOME'],
|
|
'{:}-full'.format(self._all_base_names[-1]))
|
|
if not nats_is_dir(archive_root):
|
|
warnings.warn('The input archive_root is None and the default '
|
|
'archive_root path ({:}) does not exist, try to use '
|
|
'self.archive_dir.'.format(archive_root))
|
|
archive_root = self.archive_dir
|
|
if archive_root is None or not nats_is_dir(archive_root):
|
|
raise ValueError('Invalid archive_root : {:}'.format(archive_root))
|
|
if index is None:
|
|
indexes = list(range(len(self)))
|
|
else:
|
|
indexes = [index]
|
|
for idx in indexes:
|
|
if not (0 <= idx < len(self.meta_archs)): # pylint: disable=superfluous-parens
|
|
raise ValueError('invalid index of {:}'.format(idx))
|
|
xfile_path = os.path.join(archive_root,
|
|
'{:06d}.{:}'.format(idx, PICKLE_EXT))
|
|
if not nats_is_file(xfile_path):
|
|
xfile_path = os.path.join(archive_root,
|
|
'{:d}.{:}'.format(idx, PICKLE_EXT))
|
|
assert nats_is_file(xfile_path), 'invalid data path : {:}'.format(
|
|
xfile_path)
|
|
xdata = pickle_load(xfile_path)
|
|
assert isinstance(xdata, dict), 'invalid format of data in {:}'.format(
|
|
xfile_path)
|
|
self.evaluated_indexes.add(idx)
|
|
hp2archres = collections.OrderedDict()
|
|
for hp_key, results in xdata.items():
|
|
hp2archres[hp_key] = ArchResults.create_from_state_dict(results)
|
|
self._avaliable_hps.add(hp_key)
|
|
self.arch2infos_dict[idx] = hp2archres
|
|
|
|
def query_index_by_arch(self, arch):
|
|
"""Query the index of an architecture in the search space.
|
|
|
|
Args:
|
|
arch: For topology search space, the input arch can be an architecture
|
|
string such as '|nor_conv_3x3~0|+|nor_conv_3x3~0|avg_pool_3x3~1|+|skip_connect~0|nor_conv_3x3~1|skip_connect~2|'; # pylint: disable=line-too-long
|
|
or an instance that has the 'tostr' function that can
|
|
generate the architecture string;
|
|
or it is directly an architecture index, in this case,
|
|
we will check whether it is valid or not.
|
|
This function will return the index.
|
|
If return -1, it means this architecture is not in the search space.
|
|
Otherwise, it will return an intenger in
|
|
[0, the-number-of-candidates-in-the-search-space).
|
|
|
|
Raises:
|
|
ValueError: If did not find the architecture in this benchmark.
|
|
|
|
Returns:
|
|
The index of the architcture in this benchmark.
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call query_index_by_arch with arch={:}'.format(
|
|
time_string(), arch))
|
|
if isinstance(arch, int):
|
|
if 0 <= arch < len(self):
|
|
return arch
|
|
else:
|
|
raise ValueError('Invalid architecture index {:} vs [{:}, {:}].'.format(
|
|
arch, 0, len(self)))
|
|
elif isinstance(arch, str):
|
|
if arch in self.archstr2index:
|
|
arch_index = self.archstr2index[arch]
|
|
else:
|
|
arch_index = -1
|
|
elif hasattr(arch, 'tostr'):
|
|
if arch.tostr() in self.archstr2index:
|
|
arch_index = self.archstr2index[arch.tostr()]
|
|
else:
|
|
arch_index = -1
|
|
else:
|
|
arch_index = -1
|
|
return arch_index
|
|
|
|
def query_by_arch(self, arch, hp):
|
|
"""Make the current version be compatible with the old NAS-Bench-201 version."""
|
|
return self.query_info_str_by_arch(arch, hp)
|
|
|
|
def _prepare_info(self, index):
|
|
"""This is a function to load the data from disk when using fast mode."""
|
|
if index not in self.arch2infos_dict:
|
|
if self.fast_mode and self.archive_dir is not None:
|
|
self.reload(self.archive_dir, index)
|
|
elif not self.fast_mode:
|
|
if self.verbose:
|
|
print('{:} Call _prepare_info with index={:} skip because it is not'
|
|
'the fast mode.'.format(time_string(), index))
|
|
else:
|
|
raise ValueError('Invalid status: fast_mode={:} and '
|
|
'archive_dir={:}'.format(
|
|
self.fast_mode, self.archive_dir))
|
|
else:
|
|
if index not in self.evaluated_indexes:
|
|
raise ValueError('The index of {:} is not in self.evaluated_indexes, '
|
|
'there must be something wrong.'.format(index))
|
|
if self.verbose:
|
|
print('{:} Call _prepare_info with index={:} skip because it is in '
|
|
'arch2infos_dict'.format(time_string(), index))
|
|
|
|
def clear_params(self, index: int, hp: Optional[Text] = None):
|
|
"""Remove the architecture's weights to save memory.
|
|
|
|
Args:
|
|
index: the index of the target architecture
|
|
hp: a flag to controll how to clear the parameters.
|
|
-- None: clear all the weights in '01'/'12'/'90', which indicates
|
|
the number of training epochs.
|
|
-- '01' or '12' or '90': clear all the weights in
|
|
arch2infos_dict[index][hp].
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call clear_params with index={:} and hp={:}'.format(
|
|
time_string(), index, hp))
|
|
if index not in self.arch2infos_dict:
|
|
warnings.warn('The {:}-th architecture is not in the benchmark data yet, '
|
|
'no need to clear params.'.format(index))
|
|
elif hp is None:
|
|
for key, result in self.arch2infos_dict[index].items():
|
|
result.clear_params()
|
|
else:
|
|
if str(hp) not in self.arch2infos_dict[index]:
|
|
raise ValueError('The {:}-th architecture only has hyper-parameters '
|
|
'of {:} instead of {:}.'.format(
|
|
index, list(self.arch2infos_dict[index].keys()),
|
|
hp))
|
|
self.arch2infos_dict[index][str(hp)].clear_params()
|
|
|
|
@abc.abstractmethod
|
|
def query_info_str_by_arch(self, arch, hp: Text = '12'):
|
|
"""This function is used to query the information of a specific architecture."""
|
|
|
|
def _query_info_str_by_arch(self,
|
|
arch,
|
|
hp: Text = '12',
|
|
print_information=None):
|
|
"""Internal function to query the information of `arch` when using `hp`."""
|
|
arch_index = self.query_index_by_arch(arch)
|
|
self._prepare_info(arch_index)
|
|
if arch_index in self.arch2infos_dict:
|
|
if hp not in self.arch2infos_dict[arch_index]:
|
|
raise ValueError('The {:}-th architecture only has hyper-parameters of '
|
|
'{:} instead of {:}.'.format(
|
|
arch_index,
|
|
list(self.arch2infos_dict[arch_index].keys()), hp))
|
|
info = self.arch2infos_dict[arch_index][hp]
|
|
strings = print_information(info, 'arch-index={:}'.format(arch_index))
|
|
return '\n'.join(strings)
|
|
else:
|
|
warnings.warn('Find this arch-index : {:}, but this arch is not '
|
|
'evaluated.'.format(arch_index))
|
|
return None
|
|
|
|
def query_meta_info_by_index(self, arch_index, hp: Text = '12'):
|
|
"""Return ArchResults for the 'arch_index'-th architecture."""
|
|
if self.verbose:
|
|
print('Call query_meta_info_by_index with arch_index={:}, hp={:}'.format(
|
|
arch_index, hp))
|
|
self._prepare_info(arch_index)
|
|
if arch_index in self.arch2infos_dict:
|
|
if str(hp) not in self.arch2infos_dict[arch_index]:
|
|
raise ValueError('The {:}-th architecture only has hyper-parameters of '
|
|
'{:} instead of {:}.'.format(
|
|
arch_index,
|
|
list(self.arch2infos_dict[arch_index].keys()),
|
|
hp))
|
|
info = self.arch2infos_dict[arch_index][str(hp)]
|
|
else:
|
|
raise ValueError('arch_index [{:}] does not in arch2infos'.format(
|
|
arch_index))
|
|
return copy.deepcopy(info)
|
|
|
|
def query_by_index(self,
|
|
arch_index: int,
|
|
dataname: Union[None, Text] = None,
|
|
hp: Text = '12'):
|
|
"""Query the information with the training of 01/12/90/200 epochs.
|
|
|
|
Args:
|
|
arch_index: The architecture index in this benchmark.
|
|
dataname: If dataname is None, return the ArchResults; otherwise, we will
|
|
return a dict with all trials on that dataset
|
|
(the key is the seed).
|
|
Options are 'cifar10-valid', 'cifar10', 'cifar100',
|
|
and 'ImageNet16-120'.
|
|
-- cifar10-valid : train the model on CIFAR-10 training set.
|
|
-- cifar10 : train the model on CIFAR-10 training + validation set.
|
|
-- cifar100 : train the model on CIFAR-100 training set.
|
|
-- ImageNet16-120 : train the model on ImageNet16-120 training set.
|
|
hp: The hyperparameters.
|
|
If hp=01, we train the model by 01 epochs.
|
|
If hp=12, we train the model by 01 epochs.
|
|
If hp=90, we train the model by 01 epochs.
|
|
If hp=200, we train the model by 01 epochs.
|
|
See github.com/D-X-Y/AutoDL-Projects/configs/nas-benchmark/hyper-opts
|
|
for more details.
|
|
|
|
Raises:
|
|
ValueError: If not find the matched serach space description.
|
|
|
|
Returns:
|
|
An instance fo ArchResults.
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call query_by_index with arch_index={:}, dataname={:}, '
|
|
'hp={:}'.format(time_string(), arch_index, dataname, hp))
|
|
info = self.query_meta_info_by_index(arch_index, str(hp))
|
|
if dataname is None:
|
|
return info
|
|
else:
|
|
if dataname not in info.get_dataset_names():
|
|
raise ValueError('invalid dataset-name : {:} vs. {:}'.format(
|
|
dataname, info.get_dataset_names()))
|
|
return info.query(dataname)
|
|
|
|
def find_best(self,
|
|
dataset,
|
|
metric_on_set,
|
|
flop_max=None,
|
|
param_max=None,
|
|
hp: Text = '12'):
|
|
"""Find the architecture with the highest accuracy based on some constraints."""
|
|
if self.verbose:
|
|
print('{:} Call find_best with dataset={:}, metric_on_set={:}, hp={:} '
|
|
'| with #FLOPs < {:} and #Params < {:}'.format(
|
|
time_string(), dataset, metric_on_set, hp, flop_max, param_max))
|
|
dataset, metric_on_set = remap_dataset_set_names(
|
|
dataset, metric_on_set, self.verbose)
|
|
best_index, highest_accuracy = -1, None
|
|
evaluated_indexes = sorted(list(self.evaluated_indexes))
|
|
for arch_index in evaluated_indexes:
|
|
self._prepare_info(arch_index)
|
|
arch_info = self.arch2infos_dict[arch_index][hp]
|
|
info = arch_info.get_compute_costs(dataset) # the information of costs
|
|
flop, param, latency = info['flops'], info['params'], info['latency']
|
|
if flop_max is not None and flop > flop_max:
|
|
continue
|
|
if param_max is not None and param > param_max:
|
|
continue
|
|
xinfo = arch_info.get_metrics(
|
|
dataset, metric_on_set) # the information of loss and accuracy
|
|
loss, accuracy = xinfo['loss'], xinfo['accuracy']
|
|
if best_index == -1:
|
|
best_index, highest_accuracy = arch_index, accuracy
|
|
elif highest_accuracy < accuracy:
|
|
best_index, highest_accuracy = arch_index, accuracy
|
|
del latency, loss
|
|
if self.verbose:
|
|
print(' the best architecture : [{:}] {:} with accuracy={:.3f}%'.format(
|
|
best_index, self.arch(best_index), highest_accuracy))
|
|
return best_index, highest_accuracy
|
|
|
|
def get_net_param(self, index, dataset, seed: Optional[int], hp: Text = '12'):
|
|
"""Obtain the trained weights of the `index`-th arch on `dataset`.
|
|
|
|
Args:
|
|
index: The architecture index.
|
|
dataset: The training dataset name.
|
|
seed:
|
|
-- None : return a dict containing the trained weights of all trials,
|
|
where each key is a seed and its corresponding value
|
|
is the weights.
|
|
-- Interger : return the weights of a specific trial, whose seed
|
|
is this interger.
|
|
hp:
|
|
-- 01 : train the model by 01 epochs
|
|
-- 12 : train the model by 12 epochs
|
|
-- 90 : train the model by 90 epochs
|
|
-- 200 : train the model by 200 epochs
|
|
Returns:
|
|
PyTorch weights.
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call the get_net_param function with index={:}, dataset={:}, '
|
|
'seed={:}, hp={:}'.format(time_string(), index, dataset, seed, hp))
|
|
info = self.query_meta_info_by_index(index, hp)
|
|
return info.get_net_param(dataset, seed)
|
|
|
|
def get_net_config(self, index: int, dataset: Text):
|
|
"""Obtain the configuration for the `index`-th architecture on `dataset`.
|
|
|
|
Args:
|
|
index: The architecture index.
|
|
dataset: 4 possible options as follows,
|
|
-- cifar10-valid : train the model on the CIFAR-10 training set.
|
|
-- cifar10 : train the model on the CIFAR-10 training + validation set.
|
|
-- cifar100 : train the model on the CIFAR-100 training set.
|
|
-- ImageNet16-120 : train the model on the ImageNet16-120 training set.
|
|
Returns:
|
|
A dict.
|
|
|
|
Note: some examlpes for using this function:
|
|
config = api.get_net_config(128, 'cifar10')
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call the get_net_config function with index={:}, '
|
|
'dataset={:}.'.format(time_string(), index, dataset))
|
|
self._prepare_info(index)
|
|
if index in self.arch2infos_dict:
|
|
info = self.arch2infos_dict[index]
|
|
else:
|
|
raise ValueError(
|
|
'The arch_index={:} is not in arch2infos_dict.'.format(index))
|
|
info = next(iter(info.values()))
|
|
results = info.query(dataset, None)
|
|
results = next(iter(results.values()))
|
|
return results.get_config(None)
|
|
|
|
def get_cost_info(self,
|
|
index: int,
|
|
dataset: Text,
|
|
hp: Text = '12') -> Dict[Text, float]:
|
|
"""To obtain the cost metric for the `index`-th architecture on a dataset."""
|
|
if self.verbose:
|
|
print('{:} Call the get_cost_info function with index={:}, '
|
|
'dataset={:}, and hp={:}.'.format(
|
|
time_string(), index, dataset, hp))
|
|
self._prepare_info(index)
|
|
info = self.query_meta_info_by_index(index, hp)
|
|
return info.get_compute_costs(dataset)
|
|
|
|
def get_latency(self, index: int, dataset: Text, hp: Text = '12') -> float:
|
|
"""Obtain the latency of the network.
|
|
|
|
Note: by default it will return the latency with the batch size of 256.
|
|
Args:
|
|
index: the index of the target architecture
|
|
dataset: the dataset name (cifar10-valid, cifar10, cifar100,
|
|
and ImageNet16-120)
|
|
hp: the hyperparamete indicator.
|
|
|
|
Returns:
|
|
return a float value in seconds
|
|
"""
|
|
if self.verbose:
|
|
print('{:} Call the get_latency function with index={:}, '
|
|
'dataset={:}, and hp={:}.'.format(
|
|
time_string(), index, dataset, hp))
|
|
cost_dict = self.get_cost_info(index, dataset, hp)
|
|
return cost_dict['latency']
|
|
|
|
@abc.abstractmethod
|
|
def show(self, index=-1):
|
|
"""This function will print the information of a specific (or all) architecture(s)."""
|
|
|
|
def _show(self, index=-1, print_information=None) -> None:
|
|
"""Print the information of a specific (or all) architecture(s).
|
|
|
|
Args:
|
|
index: If the index < 0: it will loop for all architectures and print
|
|
their information one by one. Else: it will print the information
|
|
of the 'index'-th architecture.
|
|
|
|
print_information: A function to print result.
|
|
|
|
Returns: None
|
|
"""
|
|
if index < 0: # show all architectures
|
|
print(self)
|
|
evaluated_indexes = sorted(list(self.evaluated_indexes))
|
|
for i, idx in enumerate(evaluated_indexes):
|
|
print('\n' + '-' * 10 + ' The ({:5d}/{:5d}) {:06d}-th '
|
|
'architecture! '.format(i, len(evaluated_indexes), idx) + '-'*10)
|
|
print('arch : {:}'.format(self.meta_archs[idx]))
|
|
for unused_key, result in self.arch2infos_dict[index].items():
|
|
strings = print_information(result)
|
|
print('>' * 40 + ' {:03d} epochs '.format(
|
|
result.get_total_epoch()) + '>' * 40)
|
|
print('\n'.join(strings))
|
|
print('<' * 40 + '------------' + '<' * 40)
|
|
else:
|
|
if 0 <= index < len(self.meta_archs):
|
|
if index not in self.evaluated_indexes:
|
|
self._prepare_info(index)
|
|
if index not in self.evaluated_indexes:
|
|
print('The {:}-th architecture has not been evaluated '
|
|
'or not saved.'.format(index))
|
|
else:
|
|
# arch_info = self.arch2infos_dict[index]
|
|
for unused_key, result in self.arch2infos_dict[index].items():
|
|
strings = print_information(result)
|
|
print('>' * 40 + ' {:03d} epochs '.format(
|
|
result.get_total_epoch()) + '>' * 40)
|
|
print('\n'.join(strings))
|
|
print('<' * 40 + '------------' + '<' * 40)
|
|
else:
|
|
print('This index ({:}) is out of range (0~{:}).'.format(
|
|
index, len(self.meta_archs)))
|
|
|
|
def statistics(self, dataset: Text, hp: Union[Text, int]) -> Dict[int, int]:
|
|
"""This function will count the number of total trials."""
|
|
if self.verbose:
|
|
print('Call the statistics function with dataset={:} and hp={:}.'.format(
|
|
dataset, hp))
|
|
valid_datasets = ['cifar10-valid', 'cifar10', 'cifar100', 'ImageNet16-120']
|
|
if dataset not in valid_datasets:
|
|
raise ValueError('{:} not in {:}'.format(dataset, valid_datasets))
|
|
nums, hp = collections.defaultdict(lambda: 0), str(hp)
|
|
# for index in range(len(self)):
|
|
for index in self.evaluated_indexes:
|
|
arch_info = self.arch2infos_dict[index][hp]
|
|
dataset_seed = arch_info.dataset_seed
|
|
if dataset not in dataset_seed:
|
|
nums[0] += 1
|
|
else:
|
|
nums[len(dataset_seed[dataset])] += 1
|
|
return dict(nums)
|
|
|
|
|
|
class ArchResults(object):
|
|
"""A class to maintain the results of an architecture under different settings."""
|
|
|
|
def __init__(self, arch_index, arch_str):
|
|
self.arch_index = int(arch_index)
|
|
self.arch_str = copy.deepcopy(arch_str)
|
|
self.all_results = dict()
|
|
self.dataset_seed = dict()
|
|
self.clear_net_done = False
|
|
|
|
def get_compute_costs(self, dataset):
|
|
"""Return the computation cost on the input dataset."""
|
|
x_seeds = self.dataset_seed[dataset]
|
|
results = [self.all_results[(dataset, seed)] for seed in x_seeds]
|
|
|
|
flops = [result.flop for result in results]
|
|
params = [result.params for result in results]
|
|
latencies = [result.get_latency() for result in results]
|
|
latencies = [x for x in latencies if x > 0]
|
|
mean_latency = np.mean(latencies) if len(latencies) else None
|
|
time_infos = collections.defaultdict(list)
|
|
for result in results:
|
|
time_info = result.get_times()
|
|
for key, value in time_info.items():
|
|
time_infos[key].append(value)
|
|
|
|
info = {
|
|
'flops': np.mean(flops),
|
|
'params': np.mean(params),
|
|
'latency': mean_latency
|
|
}
|
|
for key, value in time_infos.items():
|
|
if len(value) and value[0] is not None:
|
|
info[key] = np.mean(value)
|
|
else:
|
|
info[key] = None
|
|
return info
|
|
|
|
def get_metrics(self, dataset, setname, iepoch=None, is_random=False):
|
|
"""Obtain the loss, accuracy, etc information on a specific dataset.
|
|
|
|
If not specify, each set refer to the proposed split in NAS-Bench-201.
|
|
If some args return None or raise error, then it is not avaliable.
|
|
========================================
|
|
|
|
Args:
|
|
dataset: 4 possible options as follows
|
|
-- cifar10-valid : train the model on the CIFAR-10 training set.
|
|
-- cifar10 : train the model on the CIFAR-10 training + validation set.
|
|
-- cifar100 : train the model on the CIFAR-100 training set.
|
|
-- ImageNet16-120 : train the model on the ImageNet16-120 training set.
|
|
setname: each dataset has different setnames
|
|
-- When dataset = cifar10-valid, you can use 'train',
|
|
'x-valid', and 'ori-test'
|
|
------ 'train' : the metric on the training set.
|
|
------ 'x-valid' : the metric on the validation set.
|
|
------ 'ori-test' : the metric on the test set.
|
|
-- When dataset = cifar10, you can use 'train', 'ori-test'.
|
|
------ 'train' : the metric on the training + validation set.
|
|
------ 'ori-test' : the metric on the test set.
|
|
-- When dataset = cifar100 or ImageNet16-120, you can use 'train',
|
|
'ori-test', 'x-valid', and 'x-test'
|
|
------ 'train' : the metric on the training set.
|
|
------ 'x-valid' : the metric on the validation set.
|
|
------ 'x-test' : the metric on the test set.
|
|
------ 'ori-test' : the metric on the validation + test set.
|
|
iepoch: (None or an integer in [0, the-number-of-total-training-epochs)
|
|
------ None : return the metric after the last training epoch.
|
|
------ an integer i : return the metric after the i-th training epoch.
|
|
is_random:
|
|
------ True : return the metric of a randomly selected trial.
|
|
------ False : return the averaged metric of all avaliable trials.
|
|
------ an integer indicating the 'seed' value : return the metric of a
|
|
specific trial (whose random seed is 'is_random').
|
|
|
|
Returns:
|
|
All the metrics given the input setting.
|
|
"""
|
|
x_seeds = self.dataset_seed[dataset]
|
|
results = [self.all_results[(dataset, seed)] for seed in x_seeds]
|
|
infos = collections.defaultdict(list)
|
|
for result in results:
|
|
if setname == 'train':
|
|
info = result.get_train(iepoch)
|
|
else:
|
|
info = result.get_eval(setname, iepoch)
|
|
for key, value in info.items():
|
|
infos[key].append(value)
|
|
return_info = dict()
|
|
if isinstance(is_random, bool) and is_random: # randomly select one
|
|
index = random.randint(0, len(results)-1)
|
|
for key, value in infos.items():
|
|
return_info[key] = value[index]
|
|
elif isinstance(is_random, bool) and not is_random: # average
|
|
for key, value in infos.items():
|
|
if len(value) and value[0] is not None:
|
|
return_info[key] = np.mean(value)
|
|
else:
|
|
return_info[key] = None
|
|
elif isinstance(is_random, int): # specify the seed
|
|
if is_random not in x_seeds:
|
|
raise ValueError('can not find random seed ({:}) from {:}'.format(
|
|
is_random, x_seeds))
|
|
index = x_seeds.index(is_random)
|
|
for key, value in infos.items():
|
|
return_info[key] = value[index]
|
|
else:
|
|
raise ValueError('invalid value for is_random: {:}'.format(is_random))
|
|
return return_info
|
|
|
|
# def show(self, is_print=False):
|
|
# return print_information(self, None, is_print)
|
|
|
|
def get_dataset_names(self):
|
|
return list(self.dataset_seed.keys())
|
|
|
|
def get_dataset_seeds(self, dataset):
|
|
return copy.deepcopy(self.dataset_seed[dataset])
|
|
|
|
def get_net_param(self, dataset: Text, seed: Union[None, int] = None):
|
|
"""Return the trained network's weights on the 'dataset'.
|
|
|
|
Args:
|
|
dataset: 'cifar10-valid', 'cifar10', 'cifar100', or 'ImageNet16-120'.
|
|
seed: an integer indicates the seed value
|
|
or None that indicates returing all trials.
|
|
|
|
Returns:
|
|
The trained weights (parameters).
|
|
"""
|
|
if seed is None:
|
|
x_seeds = self.dataset_seed[dataset]
|
|
return {
|
|
seed: self.all_results[(dataset, seed)].get_net_param()
|
|
for seed in x_seeds
|
|
}
|
|
else:
|
|
xkey = (dataset, seed)
|
|
if xkey in self.all_results:
|
|
return self.all_results[xkey].get_net_param()
|
|
else:
|
|
raise ValueError('key={:} not in {:}'.format(
|
|
xkey, list(self.all_results.keys())))
|
|
|
|
def reset_latency(self, dataset: Text, seed: Union[None, Text],
|
|
latency: float) -> None:
|
|
"""This function is used to reset the latency in all corresponding ResultsCount(s)."""
|
|
if seed is None:
|
|
for seed in self.dataset_seed[dataset]:
|
|
self.all_results[(dataset, seed)].update_latency([latency])
|
|
else:
|
|
self.all_results[(dataset, seed)].update_latency([latency])
|
|
|
|
def reset_pseudo_train_times(self, dataset: Text, seed: Union[None, Text],
|
|
estimated_per_epoch_time: float) -> None:
|
|
"""This function is used to reset the train-times in all corresponding ResultsCount(s)."""
|
|
if seed is None:
|
|
for seed in self.dataset_seed[dataset]:
|
|
self.all_results[(
|
|
dataset, seed)].reset_pseudo_train_times(estimated_per_epoch_time)
|
|
else:
|
|
self.all_results[(
|
|
dataset, seed)].reset_pseudo_train_times(estimated_per_epoch_time)
|
|
|
|
def reset_pseudo_eval_times(self, dataset: Text, seed: Union[None, Text],
|
|
eval_name: Text,
|
|
estimated_per_epoch_time: float) -> None:
|
|
"""This function is used to reset the eval-times in all corresponding ResultsCount(s)."""
|
|
if seed is None:
|
|
for seed in self.dataset_seed[dataset]:
|
|
self.all_results[(dataset, seed)].reset_pseudo_eval_times(
|
|
eval_name, estimated_per_epoch_time)
|
|
else:
|
|
self.all_results[(dataset, seed)].reset_pseudo_eval_times(
|
|
eval_name, estimated_per_epoch_time)
|
|
|
|
def get_latency(self, dataset: Text) -> float:
|
|
"""Get the latency of a model on the target dataset."""
|
|
latencies = []
|
|
for seed in self.dataset_seed[dataset]:
|
|
latency = self.all_results[(dataset, seed)].get_latency()
|
|
if not isinstance(latency, float) or latency <= 0:
|
|
raise ValueError('invalid latency of {:} with seed={:} : {:}'.format(
|
|
dataset, seed, latency))
|
|
latencies.append(latency)
|
|
return sum(latencies) / len(latencies)
|
|
|
|
def get_total_epoch(self, dataset=None):
|
|
"""Return the total number of training epochs."""
|
|
if dataset is None:
|
|
epochss = []
|
|
for xdata, x_seeds in self.dataset_seed.items():
|
|
epochss += [
|
|
self.all_results[(xdata, seed)].get_total_epoch()
|
|
for seed in x_seeds
|
|
]
|
|
elif isinstance(dataset, str):
|
|
x_seeds = self.dataset_seed[dataset]
|
|
epochss = [
|
|
self.all_results[(dataset, seed)].get_total_epoch()
|
|
for seed in x_seeds
|
|
]
|
|
else:
|
|
raise ValueError('invalid dataset={:}'.format(dataset))
|
|
if len(set(epochss)) > 1:
|
|
raise ValueError(
|
|
'Each trial mush have the same number of training epochs : {:}'
|
|
.format(epochss))
|
|
return epochss[-1]
|
|
|
|
def query(self, dataset, seed=None):
|
|
"""Return the ResultsCount object (containing all information of a single trial) for 'dataset' and 'seed'."""
|
|
if seed is None:
|
|
x_seeds = self.dataset_seed[dataset]
|
|
return {seed: self.all_results[(dataset, seed)] for seed in x_seeds}
|
|
else:
|
|
return self.all_results[(dataset, seed)]
|
|
|
|
def arch_idx_str(self):
|
|
return '{:06d}'.format(self.arch_index)
|
|
|
|
def update(self, dataset_name, seed, result):
|
|
"""Update the result for the given dataset and seed."""
|
|
if dataset_name not in self.dataset_seed:
|
|
self.dataset_seed[dataset_name] = []
|
|
if seed in self.dataset_seed[dataset_name]:
|
|
raise ValueError('{:}-th arch alreadly has this seed ({:}) on {:}'.format(
|
|
self.arch_index, seed, dataset_name))
|
|
self.dataset_seed[dataset_name].append(seed)
|
|
self.dataset_seed[dataset_name] = sorted(self.dataset_seed[dataset_name])
|
|
assert (dataset_name, seed) not in self.all_results
|
|
self.all_results[(dataset_name, seed)] = result
|
|
self.clear_net_done = False
|
|
|
|
def state_dict(self):
|
|
"""Return a dict that can be used to re-create this instance."""
|
|
state_dict = dict()
|
|
for key, value in self.__dict__.items():
|
|
if key == 'all_results': # contain the class of ResultsCount
|
|
xvalue = dict()
|
|
if not isinstance(value, dict):
|
|
raise ValueError('invalid type of value for {:} : {:}'.format(
|
|
key, type(value)))
|
|
for cur_k, cur_v in value.items():
|
|
if not isinstance(cur_v, ResultsCount):
|
|
raise ValueError('invalid type of value for {:}/{:} : {:}'.format(
|
|
key, cur_k, type(cur_v)))
|
|
xvalue[cur_k] = cur_v.state_dict()
|
|
else:
|
|
xvalue = value
|
|
state_dict[key] = xvalue
|
|
return state_dict
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Update self based on the input dict."""
|
|
new_state_dict = dict()
|
|
for key, value in state_dict.items():
|
|
if key == 'all_results': # To convert to the class of ResultsCount
|
|
xvalue = dict()
|
|
if not isinstance(value, dict):
|
|
raise ValueError('invalid type of value for {:} : {:}'.format(
|
|
key, type(value)))
|
|
for cur_k, cur_v in value.items():
|
|
xvalue[cur_k] = ResultsCount.create_from_state_dict(cur_v)
|
|
else: xvalue = value
|
|
new_state_dict[key] = xvalue
|
|
self.__dict__.update(new_state_dict)
|
|
|
|
@staticmethod
|
|
def create_from_state_dict(state_dict_or_file):
|
|
"""Create the ArchResults instance from a dict or a file."""
|
|
x = ArchResults(-1, -1)
|
|
if isinstance(state_dict_or_file, str): # a file path
|
|
state_dict = pickle_load(state_dict_or_file)
|
|
elif isinstance(state_dict_or_file, dict):
|
|
state_dict = state_dict_or_file
|
|
else:
|
|
raise ValueError('invalid type of state_dict_or_file : {:}'.format(
|
|
type(state_dict_or_file)))
|
|
x.load_state_dict(state_dict)
|
|
return x
|
|
|
|
def clear_params(self):
|
|
"""Clear the weights saved in each 'result'."""
|
|
# NOTE(xuanyidong): This can help reduce the memory footprint.
|
|
for unused_key, result in self.all_results.items():
|
|
del result.net_state_dict
|
|
result.net_state_dict = None
|
|
self.clear_net_done = True
|
|
|
|
def debug_test(self):
|
|
"""Help debug and test, which will call most methods."""
|
|
all_dataset = ['cifar10-valid', 'cifar10', 'cifar100', 'ImageNet16-120']
|
|
for dataset in all_dataset:
|
|
print('---->>>> {:}'.format(dataset))
|
|
print('The latency on {:} is {:} s'.format(
|
|
dataset, self.get_latency(dataset)))
|
|
for seed in self.dataset_seed[dataset]:
|
|
result = self.all_results[(dataset, seed)]
|
|
print(' ==>> result = {:}'.format(result))
|
|
print(' ==>> cost = {:}'.format(result.get_times()))
|
|
|
|
def __repr__(self):
|
|
return ('{name}(arch-index={index}, arch={arch}, '
|
|
'{num} runs, clear={clear})'.format(
|
|
name=self.__class__.__name__,
|
|
index=self.arch_index,
|
|
arch=self.arch_str,
|
|
num=len(self.all_results),
|
|
clear=self.clear_net_done))
|
|
|
|
|
|
class ResultsCount(object):
|
|
"""ResultsCount is to save the information of one trial for a single architecture."""
|
|
|
|
def __init__(self, name, state_dict, train_accs, train_losses, params, flop,
|
|
arch_config, seed, epochs, latency):
|
|
self.name = name
|
|
self.net_state_dict = state_dict
|
|
self.train_acc1es = copy.deepcopy(train_accs)
|
|
self.train_acc5es = None
|
|
self.train_losses = copy.deepcopy(train_losses)
|
|
self.train_times = None
|
|
self.arch_config = copy.deepcopy(arch_config)
|
|
self.params = params
|
|
self.flop = flop
|
|
self.seed = seed
|
|
self.epochs = epochs
|
|
self.latency = latency
|
|
# evaluation results
|
|
self.reset_eval()
|
|
|
|
def update_train_info(self, train_acc1es, train_acc5es, train_losses,
|
|
train_times) -> None:
|
|
self.train_acc1es = train_acc1es
|
|
self.train_acc5es = train_acc5es
|
|
self.train_losses = train_losses
|
|
self.train_times = train_times
|
|
|
|
def reset_pseudo_train_times(self, estimated_per_epoch_time: float) -> None:
|
|
"""Assign the training times."""
|
|
train_times = collections.OrderedDict()
|
|
for i in range(self.epochs):
|
|
train_times[i] = estimated_per_epoch_time
|
|
self.train_times = train_times
|
|
|
|
def reset_pseudo_eval_times(
|
|
self, eval_name: Text, estimated_per_epoch_time: float) -> None:
|
|
"""Assign the evaluation times."""
|
|
if eval_name not in self.eval_names:
|
|
raise ValueError('invalid eval name : {:}'.format(eval_name))
|
|
for i in range(self.epochs):
|
|
self.eval_times['{:}@{:}'.format(eval_name, i)] = estimated_per_epoch_time
|
|
|
|
def reset_eval(self):
|
|
self.eval_names = []
|
|
self.eval_acc1es = {}
|
|
self.eval_times = {}
|
|
self.eval_losses = {}
|
|
|
|
def update_latency(self, latency):
|
|
self.latency = copy.deepcopy(latency)
|
|
|
|
def get_latency(self) -> float:
|
|
"""Return the latency value in seconds."""
|
|
# NOTE(xuanyidong): -1 represents not avaliable,
|
|
# NOTE(xuanyidong): otherwise it should be a float value.
|
|
if self.latency is None:
|
|
return -1.0
|
|
else:
|
|
return sum(self.latency) / len(self.latency)
|
|
|
|
def update_eval(self, accs, losses, times):
|
|
"""To update the evaluataion results."""
|
|
data_names = set([x.split('@')[0] for x in accs.keys()])
|
|
for data_name in data_names:
|
|
if data_name in self.eval_names:
|
|
raise ValueError('{:} has already been added into '
|
|
'eval-names'.format(data_name))
|
|
self.eval_names.append(data_name)
|
|
for iepoch in range(self.epochs):
|
|
xkey = '{:}@{:}'.format(data_name, iepoch)
|
|
self.eval_acc1es[xkey] = accs[xkey]
|
|
self.eval_losses[xkey] = losses[xkey]
|
|
self.eval_times[xkey] = times[xkey]
|
|
|
|
def update_OLD_eval(self, name, accs, losses): # pylint: disable=invalid-name
|
|
"""To update the evaluataion results (old NAS-Bench-201 version)."""
|
|
assert name not in self.eval_names, '{:} has already added'.format(name)
|
|
self.eval_names.append(name)
|
|
for iepoch in range(self.epochs):
|
|
if iepoch in accs:
|
|
self.eval_acc1es['{:}@{:}'.format(name, iepoch)] = accs[iepoch]
|
|
self.eval_losses['{:}@{:}'.format(name, iepoch)] = losses[iepoch]
|
|
|
|
def __repr__(self):
|
|
num_eval = len(self.eval_names)
|
|
set_name = '[' + ', '.join(self.eval_names) + ']'
|
|
return ('{name}({xname}, arch={arch}, FLOP={flop:.2f}M, '
|
|
'Param={param:.3f}MB, seed={seed}, {num_eval} eval-sets: '
|
|
'{set_name})'.format(name=self.__class__.__name__, xname=self.name,
|
|
arch=self.arch_config['arch_str'],
|
|
flop=self.flop, param=self.params,
|
|
seed=self.seed, num_eval=num_eval,
|
|
set_name=set_name))
|
|
|
|
def get_total_epoch(self):
|
|
return copy.deepcopy(self.epochs)
|
|
|
|
def get_times(self):
|
|
"""Obtain the information regarding both training and evaluation time."""
|
|
if self.train_times is not None and isinstance(self.train_times, dict):
|
|
train_times = list(self.train_times.values())
|
|
time_info = {
|
|
'T-train@epoch': np.mean(train_times),
|
|
'T-train@total': np.sum(train_times)
|
|
}
|
|
else:
|
|
time_info = {'T-train@epoch': None, 'T-train@total': None}
|
|
for name in self.eval_names:
|
|
try:
|
|
xtimes = [
|
|
self.eval_times['{:}@{:}'.format(name, i)]
|
|
for i in range(self.epochs)
|
|
]
|
|
time_info['T-{:}@epoch'.format(name)] = np.mean(xtimes)
|
|
time_info['T-{:}@total'.format(name)] = np.sum(xtimes)
|
|
except Exception as unused_e: # pylint: disable=broad-except
|
|
time_info['T-{:}@epoch'.format(name)] = None
|
|
time_info['T-{:}@total'.format(name)] = None
|
|
return time_info
|
|
|
|
def get_eval_set(self):
|
|
return self.eval_names
|
|
|
|
def judge_valid(self, iepoch):
|
|
if iepoch < 0 or iepoch >= self.epochs:
|
|
raise ValueError('invalid iepoch={:} < {:}'.format(iepoch, self.epochs))
|
|
|
|
def get_train(self, iepoch=None):
|
|
"""Get the training information."""
|
|
if iepoch is None: iepoch = self.epochs-1
|
|
self.judge_valid(iepoch)
|
|
if self.train_times is not None:
|
|
xtime = self.train_times[iepoch]
|
|
atime = sum([self.train_times[i] for i in range(iepoch+1)])
|
|
else:
|
|
xtime, atime = None, None
|
|
return {
|
|
'iepoch': iepoch,
|
|
'loss': self.train_losses[iepoch],
|
|
'accuracy': self.train_acc1es[iepoch],
|
|
'cur_time': xtime,
|
|
'all_time': atime
|
|
}
|
|
|
|
def get_eval(self, name, iepoch=None):
|
|
"""Get the evaluation information ; there could be multiple evaluation sets (identified by the 'name' argument)."""
|
|
if iepoch is None:
|
|
iepoch = self.epochs-1
|
|
self.judge_valid(iepoch)
|
|
|
|
def _internal_query(xname):
|
|
if isinstance(self.eval_times, dict) and len(self.eval_times):
|
|
xtime = self.eval_times['{:}@{:}'.format(xname, iepoch)]
|
|
atime = sum([
|
|
self.eval_times['{:}@{:}'.format(xname, i)]
|
|
for i in range(iepoch + 1)
|
|
])
|
|
else:
|
|
xtime, atime = None, None
|
|
return {
|
|
'iepoch': iepoch,
|
|
'loss': self.eval_losses['{:}@{:}'.format(xname, iepoch)],
|
|
'accuracy': self.eval_acc1es['{:}@{:}'.format(xname, iepoch)],
|
|
'cur_time': xtime,
|
|
'all_time': atime
|
|
}
|
|
|
|
if name == 'valid':
|
|
return _internal_query('x-valid')
|
|
else:
|
|
return _internal_query(name)
|
|
|
|
def get_net_param(self, clone=False):
|
|
if clone:
|
|
return copy.deepcopy(self.net_state_dict)
|
|
else:
|
|
return self.net_state_dict
|
|
|
|
def get_config(self, str2structure):
|
|
"""This function is used to obtain the config dict for this architecture."""
|
|
if str2structure is None:
|
|
# In this case, this is an arch in size search space of NATS-BENCH.
|
|
if 'name' in self.arch_config and self.arch_config[
|
|
'name'] == 'infer.shape.tiny':
|
|
return {
|
|
'name': 'infer.shape.tiny',
|
|
'channels': self.arch_config['channels'],
|
|
'genotype': self.arch_config['genotype'],
|
|
'num_classes': self.arch_config['class_num']
|
|
}
|
|
else: # This is an arch in NATS-BENCH's topology search space.
|
|
return {
|
|
'name': 'infer.tiny',
|
|
'C': self.arch_config['channel'],
|
|
'N': self.arch_config['num_cells'],
|
|
'arch_str': self.arch_config['arch_str'],
|
|
'num_classes': self.arch_config['class_num']
|
|
}
|
|
else: # This is an arch in the size search space of NATS-BENCH.
|
|
if 'name' in self.arch_config and self.arch_config[
|
|
'name'] == 'infer.shape.tiny':
|
|
return {
|
|
'name': 'infer.shape.tiny',
|
|
'channels': self.arch_config['channels'],
|
|
'genotype': str2structure(self.arch_config['genotype']),
|
|
'num_classes': self.arch_config['class_num']
|
|
}
|
|
else: # This is an arch in the topology search space of NATS-BENCH.
|
|
return {
|
|
'name': 'infer.tiny',
|
|
'C': self.arch_config['channel'],
|
|
'N': self.arch_config['num_cells'],
|
|
'genotype': str2structure(self.arch_config['arch_str']),
|
|
'num_classes': self.arch_config['class_num']
|
|
}
|
|
|
|
def state_dict(self):
|
|
collected_state_dict = {key: value for key, value in self.__dict__.items()}
|
|
return collected_state_dict
|
|
|
|
def load_state_dict(self, state_dict):
|
|
self.__dict__.update(state_dict)
|
|
|
|
@staticmethod
|
|
def create_from_state_dict(state_dict):
|
|
x = ResultsCount(None, None, None, None, None, None, None, None, None, None)
|
|
x.load_state_dict(state_dict)
|
|
return x
|