Neural Architecture Search

Architecture Module

class archai.supergraph.nas.arch_module.ArchModule

ArchModule enahnces nn.Module by making a clear separation between regular weights and the architecture weights. The architecture parameters can be added using create_arch_params() method and then accessed using arch_params() method.

create_arch_params(named_params: Iterable[Tuple[str, Parameter | ParameterDict | ParameterList]]) → None

set_arch_params(arch_params: ArchParams) → None

arch_params(recurse=False, only_owned=False) → ArchParams

all_owned() → ArchParams

nonarch_params(recurse: bool) → Iterator[Parameter]

training: bool

Architecture Parameters

class archai.supergraph.nas.arch_params.ArchParams(arch_params: Iterable[Tuple[str, Parameter | ParameterDict | ParameterList]], registrar: Module | None = None)

This class holds set of learnable architecture parameter(s) for a given module. For example, one instance of this class would hold alphas for one instance of MixedOp. For sharing parameters, instance of this class can be passed around. Different algorithms may add learnable parameters for their need.

param_by_kind(kind: str | None) → Iterator[Parameter]

paramlist_by_kind(kind: str | None) → Iterator[ParameterList]

paramdict_by_kind(kind: str | None) → Iterator[ParameterDict]

has_kind(kind: str) → bool

static from_module(module: Module, recurse: bool = False) → ArchParams

static nonarch_from_module(module: Module, recurse: bool = False) → Iterator[Parameter]

static empty() → ArchParams

Architecture Trainer

class archai.supergraph.nas.arch_trainer.ArchTrainer(conf_train: Config, model: Model, checkpoint: CheckPoint | None)

compute_loss(lossfn: Callable, y: Tensor, logits: Tensor, aux_weight: float, aux_logits: Tensor | None) → Tensor

post_epoch(data_loaders: DataLoaders) → None

Cell

class archai.supergraph.nas.cell.Cell(desc: CellDesc, affine: bool, droppath: bool, trainables_from: Cell | None)

ops() → Iterable[Op]

forward(s0, s1)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

DAG Edge

class archai.supergraph.nas.dag_edge.DagEdge(desc: EdgeDesc, affine: bool, droppath: bool, template_edge: DagEdge | None)

forward(inputs: List[Tensor])

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

op() → Op

training: bool

Evaluater

class archai.supergraph.nas.evaluater.EvalResult(train_metrics: Metrics)

class archai.supergraph.nas.evaluater.Evaluater

evaluate(conf_eval: Config, model_desc_builder: ModelDescBuilder) → EvalResult

train_model(conf_train: Config, model: Module, checkpoint: CheckPoint | None) → Metrics

get_data(conf_loader: Config) → DataLoaders

create_model(conf_eval: Config, model_desc_builder: ModelDescBuilder, final_desc_filename=None, full_desc_filename=None) → Module

model_from_desc(model_desc) → Model

Experiment Runner

class archai.supergraph.nas.exp_runner.ExperimentRunner(config_filename: str, base_name: str, clean_expdir=False)

run_search(conf_search: Config) → SearchResult

run_eval(conf_eval: Config) → EvalResult

run(search=True, eval=True) → Tuple[SearchResult | None, EvalResult | None]

copy_search_to_eval() → None

model_desc_builder() → ModelDescBuilder | None

searcher() → Searcher

evaluater() → Evaluater

abstract trainer_class() → Type[ArchTrainer] | None

finalizers() → Finalizers

get_expname(is_search_or_eval: bool) → str

get_conf(is_search_or_eval: bool) → Config

Finalizers

class archai.supergraph.nas.finalizers.Finalizers

Provides base algorithms for finalizing model, cell and edge which can be overriden

For op-level finalize, just put logic in op’s finalize.

For model/cell/edge level finalize, you can override the methods in this class to customize the behavior. To override any of these methods, simply create new class in your algos folder, for example, diversity/diversity_finalizers.py. In this file create class that derives from Finalizers. Then in your algos exp_runner.py, return instance of that class in its finalizers() method.

finalize_model(model: Model, to_cpu=True, restore_device=True) → ModelDesc

finalize_cells(model: Model) → List[CellDesc]

finalize_cell(cell: Cell, cell_index: int, model_desc: ModelDesc, *args, **kwargs) → CellDesc

finalize_node(node: ModuleList, node_index: int, node_desc: NodeDesc, max_final_edges: int, *args, **kwargs) → NodeDesc

select_edges(edge_desc_ranks: List[Tuple[EdgeDesc, float]], max_final_edges: int) → List[EdgeDesc]

get_edge_ranks(node: ModuleList) → Tuple[List[EdgeDesc], List[Tuple[EdgeDesc, float]]]

finalize_edge(edge) → Tuple[EdgeDesc, float | None]

Model

class archai.supergraph.nas.model.Model(model_desc: ModelDesc, droppath: bool, affine: bool)

summary() → dict

ops() → Iterable[Op]

forward(x) → Tuple[Tensor, Tensor | None]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

device_type() → str

drop_path_prob(p: float)

Set drop path probability.

This will be called externally so any DropPath_ modules get new probability. Typically, every epoch we will reduce this probability.

training: bool

class archai.supergraph.nas.model.AuxTower(aux_tower_desc: AuxTowerDesc)

forward(x: Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

Model Description

Note: All classes in this file needs to be deepcopy compatible because

descs are used as template to create copies by macro builder.

class archai.supergraph.nas.model_desc.ConvMacroParams(ch_in: int, ch_out: int)

Holds parameters that may be altered by macro architecture

clone() → ConvMacroParams

class archai.supergraph.nas.model_desc.OpDesc(name: str, params: dict, in_len: int, trainables: Mapping | None, children: List[OpDesc] | None = None, children_ins: List[int] | None = None)

Op description that is in each edge

clone(clone_trainables=True) → OpDesc

clear_trainables() → None

state_dict() → dict

load_state_dict(state_dict) → None

class archai.supergraph.nas.model_desc.EdgeDesc(op_desc: OpDesc, input_ids: List[int])

Edge description between two nodes in the cell

clone(conv_params: ConvMacroParams | None, clear_trainables: bool) → EdgeDesc

clear_trainables() → None

state_dict() → dict

load_state_dict(state_dict) → None

class archai.supergraph.nas.model_desc.NodeDesc(edges: List[EdgeDesc], conv_params: ConvMacroParams)

clone()

clear_trainables() → None

state_dict() → dict

load_state_dict(state_dict) → None

class archai.supergraph.nas.model_desc.AuxTowerDesc(ch_in: int, n_classes: int, stride: int)

class archai.supergraph.nas.model_desc.CellType(value)

An enumeration.

Regular = 'regular'

Reduction = 'reduction'

class archai.supergraph.nas.model_desc.CellDesc(id: int, cell_type: CellType, conf_cell: Config, stems: List[OpDesc], stem_shapes: List[List[int | float]], nodes: List[NodeDesc], node_shapes: List[List[int | float]], post_op: OpDesc, out_shape: List[int | float], trainables_from: int)

clone(id: int) → CellDesc

clear_trainables() → None

state_dict() → dict

load_state_dict(state_dict) → None

reset_nodes(nodes: List[NodeDesc], node_shapes: List[List[int | float]], post_op: OpDesc, out_shape: List[int | float]) → None

nodes() → List[NodeDesc]

all_empty() → bool

all_full() → bool

class archai.supergraph.nas.model_desc.ModelDesc(conf_model_desc: Config, model_stems: List[OpDesc], pool_op: OpDesc, cell_descs: List[CellDesc], aux_tower_descs: List[AuxTowerDesc | None], logits_op: OpDesc)

reset_cells(cell_descs: List[CellDesc], aux_tower_descs: List[AuxTowerDesc | None]) → None

clear_trainables() → None

cell_descs() → List[CellDesc]

cell_type_count(cell_type: CellType) → int

clone() → ModelDesc

has_aux_tower() → bool

all_empty() → bool

all_full() → bool

state_dict() → dict

load_state_dict(state_dict) → None

save(filename: str, save_trainables=False) → str | None

static load(filename: str, load_trainables=False) → ModelDesc

Model Description Builder

class archai.supergraph.nas.model_desc_builder.ModelDescBuilder

get_reduction_indices(conf_model_desc: Config) → List[int]

Returns cell indices which reduces HxW and doubles channels

get_node_channels(conf_model_desc: Config) → List[List[int]]

Returns array of channels for each node in each cell. All nodes are assumed to have same output channels as input channels.

get_conf_cell() → Config

get_conf_dataset() → Config

get_conf_model_stems() → Config

build(conf_model_desc: Config, template: ModelDesc | None = None) → ModelDesc

main entry point for the class

build_cells(in_shapes: List[List[List[int | float]]], conf_model_desc: Config) → Tuple[List[CellDesc], List[AuxTowerDesc | None]]

get_node_count(cell_index: int) → int

build_cell(in_shapes: List[List[List[int | float]]], conf_cell: Config, cell_index: int) → CellDesc

get_trainables_from(cell_index: int) → int

get_ch(shape: List[int | float]) → int

build_cell_stems(in_shapes: List[List[List[int | float]]], conf_cell: Config, cell_index: int) → Tuple[List[List[int | float]], List[OpDesc]]

build_nodes_from_template(stem_shapes: List[List[int | float]], conf_cell: Config, cell_index: int) → Tuple[List[List[int | float]], List[NodeDesc]]

build_nodes(stem_shapes: List[List[int | float]], conf_cell: Config, cell_index: int, cell_type: CellType, node_count: int, in_shape: List[int | float], out_shape: List[int | float]) → Tuple[List[List[int | float]], List[NodeDesc]]

create_cell_templates(template: ModelDesc | None) → List[CellDesc | None]

build_model_pool(in_shapes: List[List[List[int | float]]], conf_model_desc: Config) → OpDesc

build_logits_op(in_shapes: List[List[List[int | float]]], conf_model_desc: Config) → OpDesc

get_cell_template(cell_index: int) → CellDesc | None

get_cell_type(cell_index: int) → CellType

build_cell_post_op(stem_shapes: List[List[int | float]], node_shapes: List[List[int | float]], conf_cell: Config, cell_index: int) → Tuple[List[int | float], OpDesc]

build_aux_tower(out_shape: List[int | float], conf_model_desc: Config, cell_index: int) → AuxTowerDesc | None

build_model_stems(in_shapes: List[List[List[int | float]]], conf_model_desc: Config) → List[OpDesc]

pre_build(conf_model_desc: Config) → None

hook for accomplishing any setup before build starts

seed_cell(model_desc: ModelDesc) → None

NAS-Based Utitilies

archai.supergraph.nas.nas_utils.checkpoint_empty(checkpoint: CheckPoint | None) → bool

archai.supergraph.nas.nas_utils.create_checkpoint(conf_checkpoint: Config, resume: bool) → CheckPoint | None

Creates checkpoint given its config. If resume is True then attempt is made to load existing checkpoint otherwise an empty checkpoint is created.

archai.supergraph.nas.nas_utils.get_model_stats(model: Model, input_tensor_shape=[1, 3, 32, 32], clone_model=True) → ModelStats

Operations

class archai.supergraph.nas.operations.Op

static create(op_desc: OpDesc, affine: bool, arch_params: ArchParams | None = None) → Op

get_trainables() → Mapping

set_trainables(state_dict) → None

static register_op(name: str, factory_fn: Callable, exists_ok=True) → None

finalize() → Tuple[OpDesc, float | None]

for trainable op, return final op and its rank

ops() → Iterator[Tuple[Op, float]]

Return contituent ops, if this op is primitive just return self

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.PoolBN(pool_type: str, op_desc: OpDesc, affine: bool)

AvgPool or MaxPool - BN

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.SkipConnect(op_desc: OpDesc, affine)

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.FacConv(op_desc: OpDesc, kernel_length: int, padding: int, affine: bool)

Factorized conv ReLU - Conv(Kx1) - Conv(1xK) - BN

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.ReLUConvBN(op_desc: OpDesc, kernel_size: int, stride: int, padding: int, affine: bool)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.ConvBNReLU(op_desc: OpDesc, kernel_size: int, stride: int, padding: int, affine: bool)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.DilConv(op_desc: OpDesc, kernel_size: int, stride: int, padding: int, dilation: int, affine: bool)

(Dilated) depthwise separable conv ReLU - (Dilated) depthwise separable - Pointwise - BN

If dilation == 2, 3x3 conv => 5x5 receptive field

5x5 conv => 9x9 receptive field

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.SepConv(op_desc: OpDesc, kernel_size: int, padding: int, affine: bool)

Depthwise separable conv DilConv(dilation=1) * 2

This is same as two DilConv stacked with dilation=1

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.Identity(op_desc: OpDesc)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.Zero(op_desc: OpDesc)

Represents no connection. Zero op can be thought of 1x1 kernel with fixed zero weight. For stride=1, it will produce output of same dimension as input but with all 0s. Now with stride of 2, it will zero out every other pixel in output.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.FactorizedReduce(op_desc: OpDesc, affine: bool)

reduce feature maps height/width by 2X while doubling channels using two 1x1 convs, each with stride=2.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.StemBase(reduction: int)

Abstract base class for model stems that enforces reduction property indicating amount of spatial map reductions performed by stem, i.e., reduction=2 for each stride=2

training: bool

class archai.supergraph.nas.operations.StemConv3x3(op_desc: OpDesc, affine: bool)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.StemConv3x3Relu(op_desc: OpDesc, affine: bool)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.StemConv3x3S4(op_desc, affine: bool)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.StemConv3x3S4S2(op_desc, affine: bool)

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.AvgPool2d7x7

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.PoolAdaptiveAvg2D

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.PoolMeanTensor

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.LinearOp(op_desc: OpDesc)

forward(x: Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.MergeOp(op_desc: OpDesc, affine: bool)

forward(states: List[Tensor])

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

can_drop_path() → bool

training: bool

class archai.supergraph.nas.operations.ConcateChannelsOp(op_desc: OpDesc, affine: bool)

forward(states: List[Tensor])

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.ProjectChannelsOp(op_desc: OpDesc, affine: bool)

forward(states: List[Tensor])

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.DropPath_(p: float = 0.0)

Replace values in tensor by 0. with probability p Ref: https://arxiv.org/abs/1605.07648

extra_repr()

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class archai.supergraph.nas.operations.MultiOp(op_desc: OpDesc, affine: bool)

forward(x: Tensor | List[Tensor]) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

Random Finalizers

class archai.supergraph.nas.random_finalizers.RandomFinalizers

finalize_node(node: ModuleList, node_index: int, node_desc: NodeDesc, max_final_edges: int, *args, **kwargs) → NodeDesc

Search Combinations

class archai.supergraph.nas.search_combinations.SearchCombinations

search(conf_search: Config, model_desc_builder: ModelDescBuilder, trainer_class: Type[ArchTrainer] | None, finalizers: Finalizers) → SearchResult

is_better_metrics(metrics1: Metrics | None, metrics2: Metrics | None) → bool

restore_checkpoint(conf_search: Config, macro_combinations) → Tuple[int, SearchResult | None]

record_checkpoint(macro_comb_i: int, best_result: SearchResult) → None

get_combinations(conf_search: Config) → Iterator[Tuple[int, int, int]]

save_trained(conf_search: Config, reductions: int, cells: int, nodes: int, model_metrics: ModelMetrics) → None

Save the model and metric info into a log file

Searcher

class archai.supergraph.nas.searcher.ModelMetrics(model: Model, metrics: Metrics)

class archai.supergraph.nas.searcher.SearchResult(model_desc: ModelDesc | None, search_metrics: Metrics | None, train_metrics: Metrics | None)

class archai.supergraph.nas.searcher.Searcher

search(conf_search: Config, model_desc_builder: ModelDescBuilder | None, trainer_class: Type[ArchTrainer] | None, finalizers: Finalizers) → SearchResult

clean_log_result(conf_search: Config, search_result: SearchResult) → None

build_model_desc(model_desc_builder: ModelDescBuilder, conf_model_desc: Config, reductions: int, cells: int, nodes: int) → ModelDesc

get_data(conf_loader: Config) → DataLoaders

finalize_model(model: Model, finalizers: Finalizers) → ModelDesc

search_model_desc(conf_search: Config, model_desc: ModelDesc, trainer_class: Type[ArchTrainer] | None, finalizers: Finalizers) → Tuple[ModelDesc, Metrics | None]

train_model_desc(model_desc: ModelDesc, conf_train: Config) → ModelMetrics | None

Train given description

Model Description Visualizer

archai.supergraph.nas.vis_model_desc.draw_model_desc(model_desc: ModelDesc, filepath: str | None = None, caption: str | None = None) → Tuple[Digraph | None, Digraph | None]

archai.supergraph.nas.vis_model_desc.draw_cell_desc(cell_desc: CellDesc, filepath: str | None = None, caption: str | None = None) → Digraph

make DAG plot and optionally save to filepath as .png