Creating an OFA Pruner - 2.5 English

Vitis AI Optimizer User Guide (UG1333)

Document ID
Release Date
2.5 English

The pruner requires two arguments:

  • The model to be pruned
  • The inputs needed by the model for inference
import torch
from pytorch_nndct import OFAPruner

inputs = torch.randn([1, 3, 224, 224], dtype=torch.float32)
pruner = OFAPruner(model, inputs)
Note: The input does not need to be real data. You can use randomly generated dummy data if it has the same shape and type as the real data.