参数剪枝
在这篇文章中,我们讲解模型优化过程中的参数剪枝(Pruning)环节。
在很多人的印象里,剪枝是一个复杂的算法,但在这里,它的实现“非常非常简单”。我们不需要凭空创造新的值,只需要删掉不需要的值。
- def prune_gaussians(mask_prune, parameters, optimizer):
- """
- Prune selected Gaussians.
- """
- if not mask_prune.any():
- return parameters, optimizer
- keep_mask = ~mask_prune
- for name, param in parameters.items():
- old_val = param.detach()
- new_val = old_val[keep_mask]
- parameters[name] = torch.nn.Parameter(new_val, requires_grad=True)
- new_optimizer = makeOptimizer(parameters)
- return parameters, new_optimizer