# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import copy
import random
from typing import List, Optional, Tuple
from overrides import overrides
from archai.common.config import Config
from archai.supergraph.nas.model_desc import (
AuxTowerDesc,
CellDesc,
CellType,
ConvMacroParams,
EdgeDesc,
NodeDesc,
OpDesc,
TensorShape,
TensorShapes,
TensorShapesList,
)
from archai.supergraph.nas.model_desc_builder import ModelDescBuilder
[docs]class RandOps:
"""Container to store (op_names, to_states) for each nodes"""
PRIMITIVES = [
'max_pool_3x3',
'avg_pool_3x3',
'skip_connect', # identity
'sep_conv_3x3',
'sep_conv_5x5',
'dil_conv_3x3',
'dil_conv_5x5',
# we don't allow none edge for random ops
# 'none' # this must be at the end so top1 doesn't choose it
]
def __init__(self, n_nodes:int, max_edges:int) -> None:
self.ops_and_ins:List[Tuple[List[str], List[int]]] = []
for i in range(n_nodes):
op_names = random.choices(RandOps.PRIMITIVES, k=max_edges)
to_states = random.sample(list(range(i+2)), k=max_edges)
self.ops_and_ins.append((op_names, to_states))
[docs]class RandomModelDescBuilder(ModelDescBuilder):
[docs] @overrides
def build_cells(self, in_shapes:TensorShapesList, conf_model_desc:Config)\
->Tuple[List[CellDesc], List[Optional[AuxTowerDesc]]]:
max_edges = conf_model_desc['num_edges_to_sample']
node_count = self.get_node_count(0)
# create two sets of random ops, one for each cell type
self._normal_ops = RandOps(node_count, max_edges)
self._reduction_ops = RandOps(node_count, max_edges)
return super().build_cells(in_shapes, conf_model_desc)
[docs] @overrides
def build_nodes(self, stem_shapes:TensorShapes, conf_cell:Config,
cell_index:int, cell_type:CellType, node_count:int,
in_shape:TensorShape, out_shape:TensorShape) \
->Tuple[TensorShapes, List[NodeDesc]]:
assert in_shape[0]==out_shape[0]
reduction = (cell_type==CellType.Reduction)
ops = self._reduction_ops if reduction else self._normal_ops
assert node_count == len(ops.ops_and_ins)
nodes:List[NodeDesc] = []
conv_params = ConvMacroParams(in_shape[0], out_shape[0])
for op_names, to_states in ops.ops_and_ins:
edges=[]
# add random edges
for op_name, to_state in zip(op_names, to_states):
op_desc = OpDesc(op_name,
params={
'conv': conv_params,
'stride': 2 if reduction and to_state < 2 else 1
}, in_len=1, trainables=None, children=None)
edge = EdgeDesc(op_desc, input_ids=[to_state])
edges.append(edge)
nodes.append(NodeDesc(edges=edges, conv_params=conv_params))
out_shapes = [copy.deepcopy(out_shape) for _ in range(node_count)]
return out_shapes, nodes
