在 Pytorch 中冻结单个权重答案

【问题标题】：Freezing Individual Weights in Pytorch在 Pytorch 中冻结单个权重
【发布时间】：2020-01-28 10:33:36
【问题描述】：

以下问题不是How to apply layer-wise learning rate in Pytorch? 的重复项，因为该问题旨在从训练中冻结张量的子集，而不是整个层。

我正在尝试彩票假设的 PyTorch 实现。

为此，我想冻结模型中为零的权重。以下是实现它的正确方法吗？

for name, p in model.named_parameters():
            if 'weight' in name:
                tensor = p.data.cpu().numpy()
                grad_tensor = p.grad.data.cpu().numpy()
                grad_tensor = np.where(tensor == 0, 0, grad_tensor)
                p.grad.data = torch.from_numpy(grad_tensor).to(device)

【问题讨论】：

How to apply layer-wise learning rate in Pytorch?的可能重复
@Salih Karagoz 这不是那个问题的重复。此问题询问有关从训练中省略张量子集的问题。

标签： python pytorch freeze pruning

【解决方案1】：

如果您在loss.backward() 和optimizer.step() 之前执行此操作，您所拥有的似乎会起作用（参考这些变量名的常见用法）。话虽如此，似乎有点令人费解。此外，如果您的权重是浮点值，那么comparing them to exactly zero is probably a bad idea，我们可以引入一个 epsilon 来解决这个问题。 IMO 以下内容比您提出的解决方案要干净一些：

# locate zero-value weights before training loop
EPS = 1e-6
locked_masks = {n: torch.abs(w) < EPS for n, w in model.named_parameters() if n.endswith('weight')}

...

for ... #training loop

    ...

    optimizer.zero_grad()
    loss.backward()
    # zero the gradients of interest
    for n, w in model.named_parameters():                                                                                                                                                                           
        if w.grad is not None and n in locked_masks:                                                                                                                                                                                   
            w.grad[locked_masks[n]] = 0 
    optimizer.step()

【讨论】：