def __init__(self, model_desc: ModelDesc, droppath: bool, affine: bool):
super().__init__()
# some of these fields are public as finalizer needs access to them
self.desc = model_desc
# TODO: support any number of stems
assert len(model_desc.model_stems
) == 2, "Model compiler currently only supports 2 stems"
stem0_op = Op.create(model_desc.model_stems[0], affine=affine)
stem1_op = Op.create(model_desc.model_stems[1], affine=affine)
self.model_stems = nn.ModuleList((stem0_op, stem1_op))
self.cells = nn.ModuleList()
self._aux_towers = nn.ModuleList()
for i, (cell_desc, aux_tower_desc) in \
enumerate(zip(model_desc.cell_descs(), model_desc.aux_tower_descs)):
self._build_cell(cell_desc, aux_tower_desc, droppath, affine)
# adaptive pooling output size to 1x1
self.pool_op = Op.create(model_desc.pool_op, affine=affine)
# since ch_p records last cell's output channels
# it indicates the input channel number
self.logits_op = Op.create(model_desc.logits_op, affine=affine)
def seed(self, model_desc: ModelDesc) -> None:
# for petridish we add one node with identity to s1
# this will be our seed model
for cell_desc in model_desc.cell_descs():
node_count = len(cell_desc.nodes())
assert node_count >= 1
first_node = cell_desc.nodes()[0]
# if there are no edges for 1st node, add identity to s1
if len(first_node.edges) == 0:
op_desc = OpDesc(
'skip_connect',
params={
'conv':
cell_desc.conv_params,
'stride':
2 if cell_desc.cell_type == CellType.Reduction else 1
},
in_len=1,
trainables=None,
children=None)
edge = EdgeDesc(op_desc, input_ids=[1])
first_node.edges.append(edge)
# remove empty nodes
new_nodes = [
n.clone() for n in cell_desc.nodes() if len(n.edges) > 0
]
if len(new_nodes) != len(cell_desc.nodes()):
cell_desc.reset_nodes(new_nodes, cell_desc.node_ch_out,
cell_desc.post_op.name)
self._ensure_nonempty_nodes(cell_desc)
def build(self,
conf_model_desc: Config,
template: Optional[ModelDesc] = None) -> ModelDesc:
"""main entry point for the class"""
self._init_build(conf_model_desc, template)
self.pre_build(conf_model_desc)
# input shape for the stem has same channels as channels in image
# -1 indicates, actual dimensions are not known
ds_ch = self.get_conf_dataset()['channels']
in_shapes = [[[ds_ch, -1, -1, -1]]]
# create model stems
model_stems = self.build_model_stems(in_shapes, conf_model_desc)
# create cell descriptions
cell_descs, aux_tower_descs = self.build_cells(in_shapes,
conf_model_desc)
model_pool_op = self.build_model_pool(in_shapes, conf_model_desc)
logits_op = self.build_logits_op(in_shapes, conf_model_desc)
return ModelDesc(conf_model_desc, model_stems, model_pool_op,
cell_descs, aux_tower_descs, logits_op)
def build(self, model_desc: ModelDesc, search_iter: int) -> None:
# if this is not the first iteration, we add new node to each cell
if search_iter > 0:
self.add_node(model_desc)
for cell_desc in model_desc.cell_descs():
self._build_cell(cell_desc, model_desc.params['gs_num_sample'])
def finalize_model(self,
model: Model,
to_cpu=True,
restore_device=True) -> ModelDesc:
# move model to CPU before finalize because each op will serialize
# its parameters and we don't want copy of these parameters hanging on GPU
original = model.device_type()
if to_cpu:
model.cpu()
# finalize will create copy of state and this can overflow GPU RAM
assert model.device_type() == 'cpu'
cell_descs = self.finalize_cells(model)
if restore_device:
model.to(original, non_blocking=True)
return ModelDesc(
conf_model_desc=model.desc.conf_model_desc,
model_stems=[op.finalize()[0] for op in model.model_stems],
pool_op=model.pool_op.finalize()[0],
cell_descs=cell_descs,
aux_tower_descs=model.desc.aux_tower_descs,
logits_op=model.logits_op.finalize()[0])
def build(self, model_desc: ModelDesc, search_iter: int) -> None:
cell_matrix = model_desc.params['cell_matrix']
vertex_ops = model_desc.params['vertex_ops']
self._cell_matrix, self._vertex_ops = model_matrix.prune(
cell_matrix, vertex_ops)
for cell_desc in model_desc.cell_descs():
self._build_cell(cell_desc)
def build(self, model_desc:ModelDesc, search_iter:int)->None:
# if this is not the first iteration, we add new node to each cell
if search_iter > 0:
self.add_node(model_desc)
conf = get_conf()
self._gs_num_sample = conf['nas']['search']['gs']['num_sample']
for cell_desc in model_desc.cell_descs():
self._build_cell(cell_desc, self._gs_num_sample)
def build(self, model_desc: ModelDesc, search_iter: int) -> None:
# create random op sets for two cell types
assert len(model_desc.cell_descs())
n_nodes = len(model_desc.cell_descs()[0].nodes())
max_edges = 2
# create two sets of random ops, one for each cell type
normal_ops, reduction_ops = RandOps(n_nodes, max_edges), RandOps(
n_nodes, max_edges)
for cell_desc in model_desc.cell_descs():
# select rand_ops for cell type
if cell_desc.cell_type == CellType.Regular:
rand_ops = normal_ops
elif cell_desc.cell_type == CellType.Reduction:
rand_ops = reduction_ops
else:
raise NotImplementedError(
f'CellType {cell_desc.cell_type} is not recognized')
self._build_cell(cell_desc, rand_ops)
def create_model(conf_eval: Config) -> nn.Module:
# region conf vars
dataset_name = conf_eval['loader']['dataset']['name']
final_desc_filename = conf_eval['final_desc_filename']
final_model_factory = conf_eval['final_model_factory']
full_desc_filename = conf_eval['full_desc_filename']
conf_model_desc = conf_eval['model_desc']
# endregion
if final_model_factory:
splitted = final_model_factory.rsplit('.', 1)
function_name = splitted[-1]
if len(splitted) > 1:
module_name = splitted[0]
else:
module_name = _default_module_name(dataset_name, function_name)
module = importlib.import_module(
module_name) if module_name else sys.modules[__name__]
function = getattr(module, function_name)
model = function()
logger.info({
'model_factory': True,
'module_name': module_name,
'function_name': function_name,
'params': ml_utils.param_size(model)
})
else:
# load model desc file to get template model
template_model_desc = ModelDesc.load(final_desc_filename)
model = nas_utils.model_from_conf(
full_desc_filename,
conf_model_desc,
affine=True,
droppath=True,
template_model_desc=template_model_desc)
logger.info({
'model_factory': False,
'cells_len': len(model.desc.cell_descs()),
'init_node_ch': conf_model_desc['init_node_ch'],
'n_cells': conf_model_desc['n_cells'],
'n_reductions': conf_model_desc['n_reductions'],
'n_nodes': conf_model_desc['n_nodes']
})
return model
def main():
parser = argparse.ArgumentParser(description='Visualize model description')
parser.add_argument('-f',
'--model-desc-file',
type=str,
default='models/final_model_desc5.yaml',
help='Model desc file')
args, extra_args = parser.parse_known_args()
model_desc_filepath = utils.full_path(args.model_desc_file)
model_desc = ModelDesc.load(model_desc_filepath)
out_file = pathlib.Path(model_desc_filepath).with_suffix('')
draw_model_desc(model_desc, str(out_file))
def create_model(self,
conf_eval: Config,
model_desc_builder: ModelDescBuilder,
final_desc_filename=None,
full_desc_filename=None) -> nn.Module:
assert model_desc_builder is not None, 'Default evaluater requires model_desc_builder'
# region conf vars
# if explicitly passed in then don't get from conf
if not final_desc_filename:
final_desc_filename = conf_eval['final_desc_filename']
full_desc_filename = conf_eval['full_desc_filename']
conf_model_desc = conf_eval['model_desc']
# endregion
# load model desc file to get template model
template_model_desc = ModelDesc.load(final_desc_filename)
model_desc = model_desc_builder.build(conf_model_desc,
template=template_model_desc)
# save desc for reference
model_desc.save(full_desc_filename)
model = self.model_from_desc(model_desc)
logger.info({
'model_factory':
False,
'cells_len':
len(model.desc.cell_descs()),
'init_node_ch':
conf_model_desc['model_stems']['init_node_ch'],
'n_cells':
conf_model_desc['n_cells'],
'n_reductions':
conf_model_desc['n_reductions'],
'n_nodes':
conf_model_desc['cell']['n_nodes']
})
return model
def _add_node(self, model_desc: ModelDesc,
model_desc_builder: ModelDescBuilder) -> None:
for ci, cell_desc in enumerate(model_desc.cell_descs()):
reduction = (cell_desc.cell_type == CellType.Reduction)
nodes = cell_desc.nodes()
# petridish must seed with one node
assert len(nodes) > 0
# input/output channels for all nodes are same
conv_params = nodes[0].conv_params
# assign input IDs to nodes, s0 and s1 have IDs 0 and 1
# however as we will be inserting new node before last one
input_ids = list(range(len(nodes) + 1))
assert len(input_ids) >= 2 # 2 stem inputs
op_desc = OpDesc('petridish_reduction_op' if reduction else 'petridish_normal_op',
params={
'conv': conv_params,
# specify strides for each input, later we will
# give this to each primitive
'_strides':[2 if reduction and j < 2 else 1 \
for j in input_ids],
}, in_len=len(input_ids), trainables=None, children=None)
edge = EdgeDesc(op_desc, input_ids=input_ids)
new_node = NodeDesc(edges=[edge], conv_params=conv_params)
nodes.insert(len(nodes) - 1, new_node)
# output shape of all nodes are same
node_shapes = cell_desc.node_shapes
new_node_shape = copy.deepcopy(node_shapes[-1])
node_shapes.insert(len(node_shapes) - 1, new_node_shape)
# post op needs rebuilding because number of inputs to it has changed so input/output channels may be different
post_op_shape, post_op_desc = model_desc_builder.build_cell_post_op(
cell_desc.stem_shapes, node_shapes, cell_desc.conf_cell, ci)
cell_desc.reset_nodes(nodes, node_shapes, post_op_desc,
post_op_shape)
def __init__(self, model_desc:ModelDesc, droppath:bool, affine:bool):
super().__init__()
# some of these fields are public as finalizer needs access to them
self.desc = model_desc
self.stem0_op = Op.create(model_desc.stem0_op, affine=affine)
self.stem1_op = Op.create(model_desc.stem1_op, affine=affine)
self.cells = nn.ModuleList()
self._aux_towers = nn.ModuleList()
for i, (cell_desc, aux_tower_desc) in \
enumerate(zip(model_desc.cell_descs(), model_desc.aux_tower_descs)):
self._build_cell(cell_desc, aux_tower_desc, droppath, affine)
# adaptive pooling output size to 1x1
self.pool_op = Op.create(model_desc.pool_op, affine=affine)
# since ch_p records last cell's output channels
# it indicates the input channel number
self.logits_op = Op.create(model_desc.logits_op, affine=affine)
# for i,cell in enumerate(self.cells):
# print(i, ml_utils.param_size(cell))
logger.info({'model_summary': self.summary()})
def build(self, model_desc: ModelDesc, search_iter: int) -> None:
for cell_desc in model_desc.cell_descs():
self._build_cell(cell_desc)
from archai.nas.model_desc import ModelDesc
from archai.common.common import common_init
from archai.nas.model import Model
from archai.petridish.petridish_micro_builder import PetridishMicroBuilder
from archai.nas.nas_utils import create_macro_desc
from archai.common.model_summary import summary
conf = common_init(
config_filepath='confs/petridish_cifar.yaml',
param_args=['--common.experiment_name', 'petridish_run2_seed42_eval'])
conf_eval = conf['nas']['eval']
conf_model_desc = conf_eval['model_desc']
conf_model_desc['n_cells'] = 14
template_model_desc = ModelDesc.load('final_model_desc.yaml')
model_desc = create_macro_desc(conf_model_desc, True, template_model_desc)
mb = PetridishMicroBuilder()
mb.register_ops()
model = Model(model_desc, droppath=False, affine=False)
#model.cuda()
summary(model, [64, 3, 32, 32])
exit(0)
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
from archai.nas.model_desc import ModelDesc
from archai.common.common import common_init
from archai.nas.model import Model
from archai.algos.petridish.petridish_model_desc_builder import PetridishModelBuilder
from archai.common.model_summary import summary
conf = common_init(config_filepath='confs/petridish_cifar.yaml',
param_args=['--common.experiment_name', 'petridish_run2_seed42_eval'])
conf_eval = conf['nas']['eval']
conf_model_desc = conf_eval['model_desc']
conf_model_desc['n_cells'] = 14
template_model_desc = ModelDesc.load('$expdir/final_model_desc.yaml')
model_builder = PetridishModelBuilder()
model_desc = model_builder.build(conf_model_desc, template=template_model_desc)
mb = PetridishModelBuilder()
model = Model(model_desc, droppath=False, affine=False)
summary(model, [64, 3, 32, 32])
exit(0)
def _train_dist(evaluater: Evaluater, conf_eval: Config,
model_desc_builder: ModelDescBuilder,
model_desc_filename: str, common_state) -> ConvexHullPoint:
"""Train given a model"""
common.init_from(common_state)
# region config vars
conf_model_desc = conf_eval['model_desc']
max_cells = conf_model_desc['n_cells']
conf_checkpoint = conf_eval['checkpoint']
resume = conf_eval['resume']
conf_petridish = conf_eval['petridish']
cell_count_scale = conf_petridish['cell_count_scale']
#endregion
#register ops as we are in different process now
model_desc_builder.pre_build(conf_model_desc)
model_filename = utils.append_to_filename(model_desc_filename,
'_model', '.pt')
full_desc_filename = utils.append_to_filename(model_desc_filename,
'_full', '.yaml')
metrics_filename = utils.append_to_filename(model_desc_filename,
'_metrics', '.yaml')
model_stats_filename = utils.append_to_filename(
model_desc_filename, '_model_stats', '.yaml')
# create checkpoint for this specific model desc by changing the config
checkpoint = None
if conf_checkpoint is not None:
conf_checkpoint['filename'] = model_filename.split(
'.')[0] + '_checkpoint.pth'
checkpoint = nas_utils.create_checkpoint(conf_checkpoint, resume)
if checkpoint is not None and resume:
if 'metrics_stats' in checkpoint:
# return the output we had recorded in the checkpoint
convex_hull_point = checkpoint['metrics_stats']
return convex_hull_point
# template model is what we used during the search
template_model_desc = ModelDesc.load(model_desc_filename)
# we first scale this model by number of cells, keeping reductions same as in search
n_cells = math.ceil(
len(template_model_desc.cell_descs()) * cell_count_scale)
n_cells = min(n_cells, max_cells)
conf_model_desc = copy.deepcopy(conf_model_desc)
conf_model_desc['n_cells'] = n_cells
conf_model_desc[
'n_reductions'] = n_reductions = template_model_desc.cell_type_count(
CellType.Reduction)
model_desc = model_desc_builder.build(conf_model_desc,
template=template_model_desc)
# save desc for reference
model_desc.save(full_desc_filename)
model = evaluater.model_from_desc(model_desc)
train_metrics = evaluater.train_model(conf_eval, model, checkpoint)
train_metrics.save(metrics_filename)
# get metrics_stats
model_stats = nas_utils.get_model_stats(model)
# save metrics_stats
with open(model_stats_filename, 'w') as f:
yaml.dump(model_stats, f)
# save model
if model_filename:
model_filename = utils.full_path(model_filename)
ml_utils.save_model(model, model_filename)
# TODO: Causes logging error at random times. Commenting out as stop-gap fix.
# logger.info({'model_save_path': model_filename})
hull_point = ConvexHullPoint(
JobStage.EVAL_TRAINED,
0,
0,
model_desc,
(n_cells, n_reductions, len(model_desc.cell_descs()[0].nodes())),
metrics=train_metrics,
model_stats=model_stats)
if checkpoint:
checkpoint.new()
checkpoint['metrics_stats'] = hull_point
checkpoint.commit()
return hull_point