Заморозка индивидуальных весов в Pytorch - PullRequest
Новая
       23

Заморозка индивидуальных весов в Pytorch

1 голос
/ 28 сентября 2019

Я пробую реализацию 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)

1 Ответ

0 голосов
/ 29 сентября 2019

То, что у вас есть, похоже, будет работать, если вы сделали это после loss.backward() и до optimizer.step() (что относится к общему использованию этих имен переменных).Тем не менее, это кажется немного запутанным.Кроме того, если ваши веса являются значениями с плавающей запятой, то сравнение их с ровно нулем, вероятно, является плохой идеей , мы могли бы ввести эпсилон, чтобы учесть это.Тем не менее, я думаю, что следующее немного более чистое решение 

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

...

for ... #training loop

    ...

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