Creating a Pruner - 1.3 English

Vitis AI Optimizer User Guide (UG1333)

Document ID
UG1333
Release Date
2020-12-17
Version
1.3 English
A pruner can be created by providing the model to be pruned and its input shape and input dtype. Note that shape is the size of the input image and does not contain batch size.
from pytorch_nndct import Pruner
from pytorch_nndct import InputSpec

pruner = Pruner(model, InputSpec(shape=(3, 224, 224), dtype=torch.float32))

For models with multiple inputs, you can use a list of InputSpec to initialize a pruner.