在 Pytorch 中冻结个体权重

Question

以下问题不是如何在 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)

Answer 1

loss.backward()如果您在之后和之前执行了该optimizer.step()操作（请参阅这些变量名称的常见用法），您所拥有的似乎会起作用。这么一说，好像有点绕。另外，如果你的权重是浮点值，那么将它们与零进行比较可能是一个坏主意，我们可以引入一个 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()