Creating a Pruner - 1.4.1 English

Vitis AI Optimizer User Guide (UG1333)

Document ID
Release Date
1.4.1 English

The Pruner class requires two arguments.

  • The model to be pruned
  • The inference inputs
Note: It is not necessary for the input to be real data. It can be randomly generated dummy data as long as it has the same shape and type as the real data.
import torch
from pytorch_nndct import Pruner

inputs = torch.randn([1, 3, 224, 224], dtype=torch.float32)
pruner = Pruner(model, inputs)

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