在稀疏模型中会移除剪枝后的模型中设为 0 的参数。有两种方法可用于生成最终剪枝模型。
使用剪枝 API
method = 'iterative' # or 'one_step'
runner = get_pruning_runner(model, input_signature, method)
slim_model = runner.prune(removal_ratio=0.2, mode='slim')
slim_model.load_state_dict(torch.load('model_slim.pth'))不使用剪枝 API
此方法常用于对已剪枝的模型进行量化，因为有时无法调用剪枝 API。
from pytorch_nndct.utils import slim
model = create_your_baseline_model()
slim_model = slim.load_state_dict(model, torch.load('model_slim.pth'))