Как разделить нейронную сеть на подсети в Pytorch?

Question

Я бы хотел разделить нейронную сеть на две подсети, используя Pytorch. Чтобы сделать вещи конкретными, рассмотрим этот образ:

В 1 у меня есть нейронная сеть 3x4x1. То, что я хочу, например, в эпоху 1, я хотел бы только обновить веса в подсети 1, то есть веса, которые появляются в подсети 2, должны быть заморожены. Опять же, в эпоху 2, я хотел бы тренировать веса, которые появляются в подсети 2. Остальные должны быть заморожены.

Как я могу это сделать?

Answer 1

Вы можете сделать это легко, если ваш su bnet является подмножеством слоев. То есть вам не нужно замораживать какие-либо частичные слои. Это все или ничего.

Для вашего примера это означало бы разделение скрытого слоя на два разных двухузловых слоя. Каждая из них будет принадлежать ровно к одной из подсетей, что возвращает нас ко всем или ничего.

После этого вы можете переключать отдельные слои, используя require_grad . Установка этого параметра в False для параметров отключит тренировку и остановит вес. Чтобы сделать это для всей модели, подмодели или модуля , вы должны l oop через model.parameters().

Например, с 3 входами, 1 выходом и теперь расколотый скрытый слой 2x2, он может выглядеть примерно так:

import torch.nn as nn
import torch.nn.functional as F

def set_grad(model, grad):
    for param in model.parameters():
        param.requires_grad = grad

class HalfFrozenModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.hid1 = torch.nn.Linear(3, 2)
        self.hid2 = torch.nn.Linear(3, 2)
        self.out = torch.nn.Linear(4, 1)

    def set_freeze(self, hid1=False, hid2=False):
        set_grad(self.hid1, not hid1)
        set_grad(self.hid2, not hid2)

    def forward(self, inp):
        hid1 = self.hid1(inp)
        hid2 = self.hid2(inp)
        hidden = torch.cat([hid1, hid2], 1)
        return self.out(F.relu(hidden))

Тогда вы можете тренировать одну или другую половину так:

model = HalfFrozenModel()
model.set_freeze(hid1=True)
# Do some training.
model.set_freeze(hid2=True)
# Do some more training.
# ...

---

Если вы случитесь чтобы использовать fastai , то есть концепция групп слоев, которая также используется для этого. Документация fastai подробно описывает, как это работает.