参数剪枝

在这篇文章中,我们讲解模型优化过程中的参数剪枝(Pruning)环节。

在很多人的印象里,剪枝是一个复杂的算法,但在这里,它的实现“非常非常简单”。我们不需要凭空创造新的值,只需要删掉不需要的值。

  1. def prune_gaussians(mask_prune, parameters, optimizer):
  2.     """
  3.     Prune selected Gaussians.
  4.     """
  5.     if not mask_prune.any():
  6.         return parameters, optimizer
  7.  
  8.     keep_mask = ~mask_prune
  9.     for name, param in parameters.items():
  10.         old_val = param.detach()
  11.         new_val = old_val[keep_mask]
  12.         parameters[name] = torch.nn.Parameter(new_val, requires_grad=True)
  13.  
  14.     new_optimizer = makeOptimizer(parameters)
  15.     return parameters, new_optimizer