sparselearning.utils

sparselearning.utils.accuracy_helper

Implements Top-k accuracy

sparselearning.utils.accuracy_helper.get_topk_accuracy(output: Tensor, target: Tensor, topk: Tuple = (1)) → List[float]

Computes the accuracy over the k top predictions for the specified values of k.

Parameters

  • output (torch.Tensor) – predicted labels

  • target (torch.Tensor) – groundtruth labels

  • topk (Tuple[int]) – k for which top-k should be evaluted

Returns

Top-k accuracies for each value of k supplied

Return type

List[int]

sparselearning.utils.layer_wise_density

Utility for layer wise density plots.

sparselearning.utils.layer_wise_density.plot(masking: Masking, mplot: <module 'matplotlib.pyplot' from '/Users/varun/miniconda3/envs/torch38/lib/python3.8/site-packages/matplotlib/pyplot.py'>) → <module ‘matplotlib.pyplot’ from ‘/Users/varun/miniconda3/envs/torch38/lib/python3.8/site-packages/matplotlib/pyplot.py’>

Plot layer wise density bar plot.

Parameters
Returns

matplotlib plot

Return type

pyplot

sparselearning.utils.layer_wise_density.plot_as_image(masking: Masking) → Array

Plot layer wise density as bar plot figure.

Parameters

masking (sparselearning.core.Masking) – Masking instance

Returns

Numpy array representing figure (H, W, 3)

Return type

np.ndarray

sparselearning.utils.layer_wise_density.wandb_bar(masking: Masking) → wandb.plot.bar.bar

Plot layer wise density as W&B bar plot.

Parameters

masking (sparselearning.core.Masking) – Masking instance

Returns

W&B bar plot

Return type

wandb.plot.bar

sparselearning.utils.model_serialization

Code taken from

maskrcnn-benchmark

https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/utils/model_serialization.py

sparselearning.utils.model_serialization.align_and_update_state_dicts(model_state_dict, loaded_state_dict)

Strategy: suppose that the models that we will create will have prefixes appended to each of its keys, for example due to an extra level of nesting that the original pre-trained weights from ImageNet won’t contain. For example, model.state_dict() might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains res2.conv1.weight. We thus want to match both parameters together. For that, we look for each model weight, look among all loaded keys if there is one that is a suffix of the current weight name, and use it if that’s the case. If multiple matches exist, take the one with longest size of the corresponding name. For example, for the same model as before, the pretrained weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case, we want to match backbone[0].body.conv1.weight to conv1.weight, and backbone[0].body.res2.conv1.weight to res2.conv1.weight.

sparselearning.utils.model_serialization.load_state_dict(model, loaded_state_dict)
sparselearning.utils.model_serialization.strip_prefix_if_present(state_dict, prefix)

sparselearning.utils.ops

sparselearning.utils.ops.random_perm(a: torch.Tensor) → torch.Tensor

Random shuffle a tensor.

Parameters

a (torch.Tensor) – input Tensor

Returns

shuffled Tensor

Return type

torch.Tensor

sparselearning.utils.smoothen_value

class sparselearning.utils.smoothen_value.AverageValue(n: int = 0)

Bases: object

Create a simple running average for a value (loss, etc).

add_value(val: float)

Add val to calculate updated smoothed value.

Parameters

val (float) – value to add

n: int = 0
class sparselearning.utils.smoothen_value.SmoothenValue(beta: float = 0.9, n: int = 0)

Bases: object

Create a exponentially smooth moving average for a value (loss, etc) using beta.

add_value(val: float) → None

Add val to calculate updated smoothed value.

Parameters

val (float) – value to add

beta: float = 0.9
n: int = 0

sparselearning.utils.tqdm_logging

class sparselearning.utils.tqdm_logging.TqdmLoggingHandler(stream=None)

Bases: logging.StreamHandler

Handler to pass tqdm outputs to the output file

emit(record)

Emit a record.

If a formatter is specified, it is used to format the record. The record is then written to the stream with a trailing newline. If exception information is present, it is formatted using traceback.print_exception and appended to the stream. If the stream has an ‘encoding’ attribute, it is used to determine how to do the output to the stream.

sparselearning.utils.train_helper

sparselearning.utils.train_helper.get_optimizer(model: nn.Module, **kwargs) → Tuple[optim, Tuple[lr_scheduler]]

Get model optimizer

Parameters

model (nn.Module) – Pytorch model

Returns

Optimizer, LR Scheduler(s)

Return type

Tuple[optim, Tuple[lr_scheduler]]

sparselearning.utils.train_helper.load_weights(model: nn.Module, optimizer: optim, mask: Masking, ckpt_dir: str, resume: bool = True) → Tuple[nn.Module, optim, Masking, int, int, float]

Load model, optimizers, mask from a checkpoint file (.pth).

Parameters

  • model (nn.Module) – Pytorch model

  • optimizer (torch.optim.Optimizer) – model optimizer

  • mask (sparselearning.core.Masking) – Masking instance

  • ckpt_dir (Path) – Checkpoint directory

  • resume (bool) – resume or not, if not do nothing

Returns

model, optimizer, mask, step, epoch, best_val_loss

Return type

Tuple[nn.Module, optim, Masking, int, int, float]

sparselearning.utils.train_helper.save_weights(model: nn.Module, optimizer: optim, mask: Masking, val_loss: float, step: int, epoch: int, ckpt_dir: str, is_min: bool = True)

Save progress.

Parameters

  • model (nn.Module) – Pytorch model

  • optimizer (torch.optim.Optimizer) – model optimizer

  • mask (sparselearning.core.Masking) – Masking instance

  • val_loss (float) – Current validation loss

  • step (int) – Current step

  • epoch (int) – Current epoch

  • ckpt_dir (Path) – Checkpoint directory

  • is_min (bool) – Whether current model achieves least val loss

sparselearning.utils.warmup_scheduler

class sparselearning.utils.warmup_scheduler.WarmUpLR(optimizer, total_iters: int, last_epoch: int = - 1)

Bases: torch.optim.lr_scheduler._LRScheduler

Warmup-training learning rate scheduler

Parameters

  • optimizer (torch.optim.Optimizer) – model optimizer (e.g. SGD)

  • total_iters (int) – warm up phase iterations

  • last_epoch (int) – Epoch to reset at (default -1, don’t reset)

get_lr()

We will use the first m batches, and set the learning rate to base_lr * m / total_iters