Utilities

Augmented Searcher

archai.supergraph.utils.augmented_searcher.search(conf)

Augmented Trainer

archai.supergraph.utils.augmented_trainer.run_epoch(conf, logger, model: Module, loader, loss_fn, optimizer, split_type: str, epoch=0, verbose=1, scheduler=None)

Runs epoch for given dataloader and model. If optimizer is supplied then backprop and model update is done as well. This can be called from test to train modes.

archai.supergraph.utils.augmented_trainer.train_and_eval(conf, val_ratio, val_fold, save_path, only_eval, reporter=None, metric='test')

Checkpoint

class archai.supergraph.utils.checkpoint.CheckPoint(conf_checkpoint: Config, load_existing: bool)

Callback based checkpoint model.

Start new checkpoint by calling new() and save it by calling commit(). This class is also dictionary. Items that needs be saved can be done so by setting key, value pairs after new(). As any dictionary key is set, checkpoint becomes dirty. On commit(), dictionary is saved and emptied. Invariant: checkpoint remains dirty until commit() is called.

load_existing() → bool

new(*kargs, **kvargs) → None

commit() → None

is_empty() → bool

subscribe(callback: Callable) → None

Heatmap

archai.supergraph.utils.heatmap.heatmap(data: array, ax: Axes | None = None, xtick_labels: List[str] | None = None, ytick_labels: List[str] | None = None, cbar_kwargs: Dict[str, Any] | None = None, cbar_label: str | None = None, fmt: str | None = '{x:.2f}', **kwargs) → None

Plot a heatmap.

Parameters:

• data – Data to plot.

• ax – Axis to plot on.

• xtick_labels – Labels for the x-axis.

• ytick_labels – Labels for the y-axis.

• cbar_kwargs – Keyword arguments to pass to the color bar.

• cbar_label – Label for the color bar.

• fmt – Format of the annotations.

Metrics

class archai.supergraph.utils.metrics.Metrics(title: str, apex: ApexUtils | None, logger_freq: int = 50)

Record top1, top5, loss metrics, track best so far.

There are 3 levels of metrics: 1. Run level - these for the one call of ‘fit’, example, best top1 2. Epoch level - these are the averages maintained top1, top5, loss 3. Step level - these are for every step in epoch

The pre_run must be called before fit call which will reset all metrics. Similarly pre_epoch will reset running averages and pre_step will reset step level metrics like average step time.

The post_step will simply update the running averages while post_epoch updates best we have seen for each epoch.

pre_run() → None

post_run(test_metrics: Metrics | None = None) → None

pre_step(x: Tensor, y: Tensor)

post_step(x: Tensor, y: Tensor, logits: Tensor, loss: Tensor, steps: int) → None

pre_epoch(lr: float = nan) → None

post_epoch(lr: float = nan, val_metrics: Metrics | None = None)

state_dict() → Mapping

load_state_dict(state_dict: dict) → None

save(filepath: str) → str | None

epochs() → int

Returns epochs recorded so far

cur_epoch() → EpochMetrics

reduce_min(val)

reduce_max(val)

reduce_sum(val)

reduce_mean(val)

is_dist() → bool

best_train_top1() → float

best_val_top1() → float

best_test_top1() → float

class archai.supergraph.utils.metrics.Accumulator

add(key, value)

add_dict(dict)

get_dict()

items()

class archai.supergraph.utils.metrics.EpochMetrics(index: int)

Stores the metrics for each epoch. Training metrics is in top1, top5 etc while validation metrics is in val_metrics

pre_step()

post_step(top1: float, top5: float, loss: float, batch: int)

pre_epoch(lr: float)

post_epoch(lr: float, val_metrics: Metrics | None)

duration()

class archai.supergraph.utils.metrics.RunMetrics

Metrics for the entire run. It mainly consist of metrics for each epoch

pre_run()

post_run(test_metrics: Metrics | None = None)

add_epoch() → EpochMetrics

cur_epoch() → EpochMetrics

best_epoch() → Tuple[EpochMetrics, EpochMetrics | None, EpochMetrics | None]

epoch_time_avg()

step_time_avg()

duration()

Multi-Optimizer

class archai.supergraph.utils.multi_optim.OptimSched(optim: Optimizer, sched: _LRScheduler | None, sched_on_epoch: bool | None)

Holds the optimizer and scheduler

class archai.supergraph.utils.multi_optim.MultiOptim

append(optim_sched: OptimSched) → None

zero_grad() → None

step() → None

epoch(epoch: int | None = None) → None

get_lr(optim_index: int, param_index: int) → float

state_dict() → dict

load_state_dict(state_dict: dict) → None

Tester

class archai.supergraph.utils.tester.Tester(conf_val: Config, model: Module, apex: ApexUtils)

test(test_dl: DataLoader) → Metrics

get_metrics() → Metrics | None

state_dict() → dict

get_device()

load_state_dict(state_dict: dict) → None

Trainer

class archai.supergraph.utils.trainer.Trainer(conf_train: Config, model: Module, checkpoint: CheckPoint | None = None)

fit(data_loaders: DataLoaders) → Metrics

create_multi_optim(train_len: int) → MultiOptim

create_optimizer(conf_optim: Config, params) → Optimizer

create_scheduler(conf_sched: Config, optim: Optimizer, steps_per_epoch: int) → Tuple[_LRScheduler | None, bool]

get_optimizer(index=0) → Optimizer

get_scheduler(index=0) → _LRScheduler | None

get_metrics() → Metrics

pre_fit(data_loaders: DataLoaders) → None

post_fit(data_loaders: DataLoaders) → None

pre_epoch(data_loaders: DataLoaders) → None

post_epoch(data_loaders: DataLoaders) → None

pre_step(x: Tensor, y: Tensor) → None

post_step(x: Tensor, y: Tensor, logits: Tensor, loss: Tensor, steps: int) → None

get_device()

restore_checkpoint() → None

epoch() → int

update_checkpoint(checkpoint: CheckPoint) → None

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