nn-pipeline

Common utility functions for training and evaluating neural networks as well as analyzing computational complexity


Keywords
Complexity, Profiler, CNN, model-complexity, neural-network
License
MIT
Install
pip install nn-pipeline==0.10

Documentation

This repository is written to simplify the common neural network pipeline

In this repo, you will find: 1) code for training and evaluation of neural networks; 2) code for layerwise analysis of computational complexity

Installation:

pip install nn-pipeline

Example for analyzing complexity:

from nnutils.cnn_complexity_analyzer.profile import profile_compute_layers

# model: torch.nn.Module
# inputs: dictionary type decribing input to the network
# Example: most CNN models on cifar10: inputs={'x': torch.randn(1,3,32,32)} 
# Keys in inputs needs to be the same as the parameter names for model's forward function

# profile results will look like:
# {
#    layer1.weight: ComputeProfile
#    ...
# }
# where the keys are the same as keys in model.state_dict()
# Attributes of ComputeProfile is found in nnutils.cnn_complexity_analyzer.utils.NNComputeModuleProfile
# Number of keys in profile results is the sum of CONV layers and FC layers in the network

profile_results, model_sparsity = profile_compute_layers(model, inputs=inputs, verbose=True)

Example for evaluating accuracy:

from nnutils.training_pipeline import accuracy_evaluator

# model: torch.nn.Module
# device: 'cpu' or 'cuda' device to evalaute the network
# data_loader: dataset for the model
# criterion: loss function for the model

top1_acc, loss = accuracy_evaluator.eval(model, device, data_loader, criterion, print_acc=True)

Example for evaluating the latency:

from nnutils.cnn_complexity_analyzer.profile import profile_compute_layers
from nnutils.training_pipeline import latency_evaluator

target_kernel = 'mkl'
target_device = 'cpu'
num_trials = 5

profile_results, _ = profile_compute_layers(model, inputs=inputs, verbose=True)

net_state = model.state_dict() # state dict for model
total_latency, layerwise_latency = latency_evaluator.evaluate_latency(
    model_state=net_state,
    inputs=profile_results,
    target_kernel=target_kernel,
    target_device=target_device,
    num_trials=num_trials,
    verbose=True
) # average latency measured over num_trials runs