Package: torch-optim
Torch optimizerTorch optimizer is a Python library for optimizing PyTorch models using techniques of neural network pruning. Neural network pruning can be formulated as an optimization problem to determine best subset of set of network's weights, i. e.:Maximize:Accuracy(model(W • M))Subject to:Resourcej(model(W • M)) <= Budgetjwhere W are model's weights, M is binary mask with size |M| = |W|, resource can be any resource we want to reduce (e. g. FLOPs, MACs, latency, model size, ...) and budget is our desired upper bound of the resource we want to reduce.Library provides several functionalities to facilitate solving given optimization problem.FeaturesObjective functionsObjectivemodule provides common interface for modelling optimization objective functions. Objective function of arbitrary optimization problem can be created by implementing the interface. Module also provides several implementations of objective function in context of neural network pruning to evaluate pruned neural net's performance and efficiency, such as accuracy or relative decrease of MACs (Multiply–accumulate operations) according to number of MACs of original unpruned net.ConstraintsConstraintmodule provides common interface for modelling optimization constraints. Constraint of an arbitrary optimization problem can be created by the interface implementation. For neural network pruning purposes, a constraint that checks validity of the pruning (i. e. no layer can contain empty weight tensor after pruning) is provided.Optimization algorithmsOptimizermodule contains common interface for an optimization algortihm implementations. Module also contains implementation of Genetic algorithm (GA) meta-heuristic. Two implementations of GA are provided: 1. to solve integer optimization problems and 2. to solve binary optimization problems. Detailed description of GA implementations can be found in the module.PruningPrunermodule provides basic functionality for structured neural network pruning. Structured pruning can be performed in different levels of granularity. For channel pruning, where individual filters / neurons are pruned, a channel pruner is provided. For module level pruning, where individual layers or blocks of the network can be pruned, library provides module pruner implementation.InstallationUse the package managerpipto install Torch optimizer.pipinstalltorch-optimUsageOne can train their own PyTorch model on arbitrary dataset and then use the library functionalities to perform structured pruning. Here is a simple example:importtorchfromtorchoptimportutilsfromtorchopt.prune.prunerimportChannelPrunerfromtorchopt.optim.optimizerimportIntegerGAOptimizerfromtorchopt.optim.objectiveimportAccuracy,Macs,ObjectiveContainerfromtorchopt.optim.constraintimportChannelConstraintfromthopimportprofile# Get your trained modelmodel=torch.load('path/to/trained/model.pth')# Get dataset, on which model was trained. Dataset should be divided to training, validation# and testing set. Validation set will be used for measuring accuracy of pruned model.train_set,val_set,test_set=get_dataset()# Define model's input shapeinput_shape=(1,3,32,32)# Specify device, on wich optimization will be performeddevice="cuda"iftorch.cuda.is_available()else"cpu"# Prunable modules are all linear and convolutional layers in the netnames=[nameforname,_inutils.prunable_modules(model)]bounds=[(0,len(module.weight)-1)for_,moduleinutils.prunable_modules(model)]pruner=ChannelPruner(names,input_shape)# Create GA optimizeroptimizer=IntegerGAOptimizer(ind_size=len(names),pop_size=100,elite_num=10,tourn_size=10,n_gen=30,mutp=0.1,mut_indp=0.05,cx_indp=0.5,bounds=bounds)sample=torch.randn(input_shape,device=device)orig_acc=utils.evaluate(model,test_data,device)orig_macs,_=profile(model,inputs=(sample,),verbose=False)# Create composed objective function to get best trade-off between model accuracy and MACs reductionacc=Accuracy(model=model,pruner=pruner,weight=1.0,val_data=val_set,orig_acc=orig_acc)macs=Macs(model=model,pruner=pruner,orig_macs=orig_macs,weight=1.0,in_shape=input_shape)objective=ObjectiveContainer(acc,macs)constraint=ChannelConstraint(model=model,pruner=pruner)# Perform optimizationsolution=optimizer.maximize(objective,constraint)# Get pruned model according to best solution found by GApruned_model=pruner.prune(model=model,mask=solution)LicenseMIT