移除遮罩以完成剪枝
在對模型進行剪枝之後，你需要在儲存或部署前完成最後處理。也就是要移除剪枝遮罩，讓模型不再追蹤哪些權重被設為 0。你將使用 prune.remove() 來完成這個步驟。
本練習屬於課程
Scalable AI Models with PyTorch Lightning
練習說明
- 完成對權重參數的剪枝。
- 使用
prune.remove()移除遮罩並保存最終權重。 - 列印模型以確認剪枝遮罩已被移除。
動手互動練習
試著完成這個範例程式碼，體驗一下這個練習。
import torch.nn.utils.prune as prune
# Before pruning
print(model)
# Finalize pruning by removing the pruning mask
prune.____(model[3], "____")
# Print model structure after pruning
print(____)