剪枝器需要两个实参:
- 要剪枝的模型
- 模型推断所需的输入
import torch
from pytorch_nndct import OFAPruner
inputs = torch.randn([1, 3, 224, 224], dtype=torch.float32)
pruner = OFAPruner(model, inputs)
注释: 输入无需采用真实数据。您可使用随机生成的虚拟数据,只要其形状和类型与真实数据相同即可。
剪枝器需要两个实参:
import torch
from pytorch_nndct import OFAPruner
inputs = torch.randn([1, 3, 224, 224], dtype=torch.float32)
pruner = OFAPruner(model, inputs)