以下问题不是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)
答案 0 :(得分: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()