11.2 nn.Module és paraméterregisztráció

11.2.1 Cél, előismeret és induló kérdés

A 11.1. Tensor és autograd-életciklus egységben a tenzor életciklusát vizsgáltuk: a levéltensor, a requires_grad flag, a .grad felhalmozódás és a gráf leválasztása. Itt azt a lépést tesszük meg, amely nélkül egy több száz paraméteres háló gyakorlatilag kezelhetetlen: a tenzorokat egyetlen modellobjektumba szervezzük, hogy az optimizer és a mentő eljárás egyetlen listából lássa, mely tenzorok tanulnak.

Induló becslés: egy $3\to4\to2$ MLP két affin rétege esetén a súly-tenzorok elemszáma $3\cdot4 + 4\cdot2=20$, a torzításoké $4+2=6$, tehát összesen $26$ skalár paraméter várható. A kérdés: hogyan garantálja az osztály, hogy pontosan ez a $26$ tenzor kap gradienst, jelenik meg az optimizerben, és menthető el egyetlen állományban?

11.2.2 A Module, a Parameter és a buffer definíciója

Definíció (nn.Module). A PyTorch nn.Module osztálya két felelősséget egyesít. Egyrészt a modell nevét, az almodulokat és a tanulható tenzorokat tároló szerkezet; másrészt a számítást definiáló forward metódus hordozója. A modult nem a forward-dal, hanem a model(x) hívással futtatjuk: a __call__ hookokat futtat, például a beágyazott modulok előre- és visszafelé horgait.

Definíció (nn.Parameter). A Parameter a torch.Tensor alosztálya, amely levéltensor, és az nn.Module attribútumaként való hozzárendelés során automatikusan a modul _parameters nyilvántartásába kerül. A Parameter(t) konstruktor requires_grad=True flaget biztosít, ezért a paraméter megjelenik az autograd gráfban, és a backward a .grad mezőt frissíti.

Definíció (buffer). A register_buffer(name, tensor, persistent=True) egy közönséges tenzort rendel a modulhoz. A buffer nem tanuló állapot: nem jelenik meg a parameters() iterációban, requires_grad=False, tehát gradienst nem kap. Ugyanakkor a modul része: a .to(device) áthelyezi, és a state_dict() menti, ha persistent=True. A buffer például futó statisztikát, normalizációs konstansot vagy pozíciókódot tárolhat.

A forward bemenete a 7.2. Softmax és többosztályos cross-entropy alfejezetben látott konvenció szerint batchelt tenzor: egy $x\in\mathbb R^{N\times d}$ bemenetre a lineáris réteg $W\in\mathbb R^{m\times d}$ súllyal és $b\in\mathbb R^m$ torzítással

$$z = xW^\top+b \in \mathbb R^{N\times m}.$$

A state_dict() egy OrderedDict[str, Tensor] objektumot ad; az almodulok kulcsai ponttal tagolt hierarchikus nevet kapnak, például fc1.weight. A betöltésnél a kulcsoknak és alakoknak egyezniük kell; a részletes protokollt a 12.2. Checkpoint és folytathatóság egységben tárgyaljuk.

A train(mode=True) és eval() hívások a modul training logikai mezőjét állítják, és rekurzívan végigmennek az almodulokon. Az nn.Dropout(p=0.5) például tanítási módban az aktív elemeket $1/(1-p)=2$-vel skálázza, kiértékelési módban viszont identitástranszformáció. A BatchNorm futó statisztikái tipikus bufferek, de azok részletes elemzése a 18.2. Batch normalization train és eval módban egységben következik.

11.2.3 Tanuló paraméter és buffer: a különbség levezetése

A különbséget pontos állításként fogalmazzuk meg, majd egy konkrét modellen ellenőrizzük.

Állítás (regisztrációs szabály). Legyen $M$ egy nn.Module leszármazott. Ha a __init__-ben egy $p$ nn.Parameter kerül a self.p attribútumba, akkor:

1. $p$ megjelenik az M.parameters() iterációban, és p.requires_grad=True;
2. a "p" kulcs szerepel az M.state_dict()-ben.

Ha egy $b$ tenzort a self.register_buffer("q", b) hívással regisztrálunk, akkor $b$ nem szerepel az M.parameters()-ben, b.requires_grad=False, de a "q" kulcs megjelenik a state_dict()-ben. Továbbá M.to(device) mindkét tenzort ugyanarra az eszközre helyezi.

Indoklás. A Parameter hozzárendelését a Module.__setattr__ észleli, és a tenzort a modul _parameters rendezett szótárába teszi; a register_buffer a _buffers szótárba ír. A parameters() rekurzívan csak a _parameters gyűjteményt járja be, a state_dict() pedig mindkettőt. Az optim.Optimizer konstruktora az M.parameters() által visszaadott referencialistát kapja, ezért a bufferre soha nem lép. Ez interfész-szabály, nem analitikus tétel; a laborban független ellenőrzéssel mutatjuk meg.

A konkrét modell legyen egy affin réteg, amely a bemeneti norma exponenciális mozgóátlagát bufferben tartja. Az input $\mathbb R^{N\times 3}$, a kimenet $\mathbb R^{N\times 2}$, a buffer skalár.

# munkakönyvtár: lab/n042
import torch
import torch.nn as nn

class WeightedStats(nn.Module):
    """Affin réteg, amely az input normájának exponenciális
    mozgóátlagát bufferben tartja."""
    def __init__(self, d_in: int, d_out: int, alpha: float = 0.9):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(d_out, d_in))
        self.bias = nn.Parameter(torch.zeros(d_out))
        self.register_buffer('ema_input_norm', torch.tensor(0.0))
        self.register_buffer('n_seen', torch.tensor(0, dtype=torch.int64))
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (N, d_in)
        z = torch.nn.functional.linear(x, self.weight, self.bias)  # (N, d_out)
        with torch.no_grad():
            mean_norm = x.norm(dim=1).mean()
            self.ema_input_norm.mul_(self.alpha).add_(mean_norm * (1.0 - self.alpha))
            self.n_seen += 1
        return z

m = WeightedStats(3, 2)
print([name for name, _ in m.named_parameters()])   # ['weight', 'bias']
print([name for name, _ in m.named_buffers()])      # ['ema_input_norm', 'n_seen']
print(sorted(m.state_dict().keys()))
# ['bias', 'ema_input_norm', 'n_seen', 'weight']
print(m.ema_input_norm.requires_grad)               # False

A forward-ban a bufferfrissítést with torch.no_grad(): blokkban végezzük. Ez nem azért kötelező, mert a buffer requires_grad=False; a blokk azt is megakadályozza, hogy a frissítő műveletek bármilyen gráfélt hozzanak létre, és így a tanítás visszafelé menete változatlan marad.

Az optimizercsatlakozás következménye: torch.optim.SGD(m.parameters(), lr=0.1) kizárólag a weight és bias tenzorokra hivatkozik, tehát a buffer értéke csak a forward explicit frissítése révén változik.

11.2.4 Kézi példa és ellenőrző labor

A kézi példa az induló becslés pontossá tétele. Egy nn.Linear(3,4) réteg súlya $\mathbb R^{4\times3}$, torzítása $\mathbb R^4$, az nn.Linear(4,2) súlya $\mathbb R^{2\times4}$, torzítása $\mathbb R^2$. A paraméterszám:

$$(4\cdot3+4)+(2\cdot4+2)=12+4+8+2=26.$$

A labor egy teljes, futtatható ellenőrző szkript, amely a paraméterszámot, az eszközt, a state_dict kulcsait és a train/eval váltást vizsgálja. A kimenet a kód alapján számított várt érték; a jelen szöveg írása során a szkript nem futott, ezért „jelölt, nem mért” státusszal közöljük, a tanulólaborban a tényleges futtatás adja a bizonyítékot.

# munkakönyvtár: lab/n042
# futtatás: python n042_module_check.py
import torch
import torch.nn as nn

torch.manual_seed(0)
DEVICE = torch.device('cpu')
DTYPE = torch.float32

class MLP342(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3, 4)
        self.fc2 = nn.Linear(4, 2)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = MLP342().to(device=DEVICE, dtype=DTYPE)

n_params = sum(p.numel() for p in model.parameters())
print(n_params)                                    # 26
print(model.state_dict()['fc1.weight'].shape)      # torch.Size([4, 3])
print(model.state_dict()['fc1.bias'].shape)        # torch.Size([4])
print(sorted(model.state_dict().keys()))
# ['fc1.bias', 'fc1.weight', 'fc2.bias', 'fc2.weight']
print(next(model.parameters()).device)             # cpu
print(next(model.parameters()).dtype)              # torch.float32

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
print([len(g['params']) for g in optimizer.param_groups])  # [4]

x = torch.ones(1, 3, device=DEVICE, dtype=DTYPE)
model.eval()
with torch.no_grad():
    y = model(x)
print(y.shape)                                     # torch.Size([1, 2])
print(model.training)                              # False

# train/eval eltérése egyetlen Dropout rétegen
drop = nn.Dropout(p=0.5)
drop.eval()
print(torch.equal(drop(x), x))                     # True
drop.train()
z = drop(x)
print(z.shape)                                     # torch.Size([1, 3])
# tanítási módban az aktív elemek 0 vagy 2; a pontos maszk seedfüggő

Az ellenőrzés öt pontja: a paraméterszám pontosan $26$; a state_dict a négy hierarchikus kulcsot tartalmazza; az összes paraméter CPU-n és float32 formátumban van; az optimizer paramétercsoportjában a négy regisztrált tenzor szerepel; a Dropout kiértékelési módban determinisztikusan azonosság, tanítási módban viszont a modell training flagje vált.

11.2.5 Feladatok

F1 (fogalmi, nehézség 1). Magyarázza el, mi a különbség egy nn.Parameter, egy register_buffer-rel regisztrált tenzor és egy sima float Python-attribútum között. Adjon egy konkrét, nem triviális ellenőrzést, amely a három esetet megkülönbözteti.

F2 (levezetési/bizonyítási, nehézség 2). Egy modell két lineáris rétegből áll: $5\to6\to3$. Vezesse le (a) a rétegenkénti tenzoralakokat, (b) az összes paraméter számát, és (c) indokolja, hogy egy register_buffer('scale', torch.ones(1)) hívás után az összes paraméterszám változatlan marad, de a state_dict() kulcsszáma eggyel nő. Minden lépésben adja meg az indexhatárokat és a shape-ellenőrzést.

F3 (implementációs, nehézség 2). Írjon GatedScale modult, amelynek egyetlen tanulható paramétere a logit_threshold skalár. A forward az $x\in\mathbb R^{N\times1}$ bemenetre a $y = x\cdot \sigma(t)$ kimenetet adja, ahol $\sigma$ a szigmoid; a modul egy pos_count nevű int64 bufferben számolja, hányszor volt $\sigma(t)>0.5$, a frissítést no_grad blokkban végezze. A megoldáshoz adja meg a teljes importot, a modellt és egy kontrollt egy $2\times1$-es bemenettel.

F4 (hibakeresési, nehézség 2). Tekintse az alábbi hibás osztályt.

class Broken(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(3, 4), nn.Linear(4, 2)]

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

Az optim = torch.optim.SGD(m.parameters(), lr=0.1) hívás után az optimizer paramétercsoportja üres. Magyarázza meg a hiba okát és következményét, javítsa a kódot, és adja meg, milyen kimenettel ellenőrzi, hogy a javítás után a négy paraméter regisztrálva van. Van-e más logikai hiba a forward-ban? Ha igen, mi az, és hogyan javítja?

F5 (kísérleti/transzfer, nehézség 3). Tervezzen ellenőrzési protokollt arra, hogy egy state_dict mentése és egy második, azonos architektúrájú folyamatban való betöltése után a kimenet bitazonos legyen. Adja meg a seed, a dtype, a device, a modul train/eval állapota és az összehasonlítás mérőszáma mellett az érvényességi korlátot is: miért nem garantál egyetlen seed minden platformon bitazonosságot?

11.2.6 Megoldások és indoklás

F1 megoldás. A Parameter a modul tanuló állapota: regisztrálódik, requires_grad=True, a parameters() felsorolja, az optimizer frissíti. A buffer nem tanuló, de szerves állapot: regisztrálódik, a parameters() kihagyja, gradienst nem kap, viszont a .to() mozgatja és a state_dict() menti. A sima float attribútum egyik nyilvántartásba sem kerül: a parameters() nem látja, a state_dict() nem menti, és a .to() sem mozgatja. Ellenőrzés: egy modulban mindhárom típust létrehozzuk; len(list(m.parameters())), a state_dict() kulcshalmaza és a .to('cpu') utáni hivatkozás különböztet.

F2 megoldás. Az első réteg súlya $W_1\in\mathbb R^{6\times5}$, torzítása $b_1\in\mathbb R^6$, a másodiké $W_2\in\mathbb R^{3\times6}$, $b_2\in\mathbb R^3$. A paraméterszám:

$$(6\cdot5+6)+(3\cdot6+3)=36+21=57.$$

A state_dict() kulcsai fc1.weight, fc1.bias, fc2.weight, fc2.bias. A register_buffer('scale', torch.ones(1)) a _buffers-be ír; a parameters() továbbra is a két réteg paramétereit iterálja, ezért az összeg $57$ marad. A state_dict() viszont a _parameters és _buffers unióját adja, tehát az úgynevezett scale kulcs hozzáadódik, és a kulcsszám $4+1=5$. Shape-ellenőrzés: a buffer alakja torch.Size([1]), a paraméterszám összesítésében nem szerepel, mert requires_grad=False.

F3 megoldás.

import torch
import torch.nn as nn

class GatedScale(nn.Module):
    def __init__(self):
        super().__init__()
        self.logit_threshold = nn.Parameter(torch.tensor(0.0))
        self.register_buffer('pos_count', torch.tensor(0, dtype=torch.int64))

    def forward(self, x):
        # x: (N, 1)
        s = torch.sigmoid(self.logit_threshold)  # skalár tenzor
        y = x * s
        with torch.no_grad():
            if s.item() > 0.5:
                self.pos_count += 1
        return y

m = GatedScale()
x = torch.tensor([[1.0], [2.0]])
with torch.no_grad():
    y = m(x)
print(y)              # [[0.5], [1.0]], mert sigmoid(0)=0.5
print(m.pos_count)    # tensor(0), mert a feltétel hamis
print(len(list(m.parameters())))  # 1
print(sorted(m.state_dict().keys()))  # ['logit_threshold', 'pos_count']

A szigmoid $s=\sigma(0)=0.5$, így a kapu a bemenet felét adja; a pos_count nem növekszik, mert a feltétel hamis.

F4 megoldás. A self.layers sima Python-lista; a Module.__setattr__ nem járja be a lista elemeit, ezért a két Linear réteg nem kerül a modul _modules nyilvántartásába. Az optimizer a m.parameters() iterációt használja, amely üres, így a gradienslépés nem érint semmit; a súlyok véletlen értékeiken maradnak. A javítás az nn.ModuleList használata:

self.layers = nn.ModuleList([nn.Linear(3, 4), nn.Linear(4, 2)])

Az ellenőrzés len(list(m.parameters())) == 4, és list(m.named_modules()) listában a két lineáris almodul szerepel. A második hiba: a forward az utolsó, $2$ kimenetű réteg után is ReLU-t alkalmaz, tehát a modell nem képes negatív logitot kiadni; a javítás:

def forward(self, x):
    x = torch.relu(self.layers[0](x))
    return self.layers[1](x)

F5 megoldásminta. A protokoll: rögzítsük a generátor seedet torch.manual_seed(0); a modellt CPU-n, float32 dtype-tal, eval() módban tartsuk; determinisztikus bemenetet torch.ones(1,3) használva rögzítsük a referencia kimenetet; torch.save(model.state_dict(), 'mlp342.pt'); egy második Python-folyamatban hozzuk létre az azonos MLP342 architektúrát, töltsük be torch.load('mlp342.pt', map_location='cpu', weights_only=True) és load_state_dict hívással; ugyanazon bemenetre a kimenetet torch.equal mérőszámmal hasonlítsuk a referenciához. Az érvényességi korlát: a seed csak az adott NumPy/PyTorch verzió, CPU-generátor és architektúra mellett reprodukál; más platformon a lebegőpontos műveletek sorrendje eltérhet, ezért a feltétlen bitazonosság nem garantált.

11.2.7 Tipikus tévedés ellenpéldával

A leggyakoribb tévedés a sima Python-listás almodulok használata. A

self.layers = [nn.Linear(3, 4), nn.Linear(4, 2)]

sorban a lista elemei nem regisztrálódnak a modulba; az m.parameters() üres. Az nn.ModuleList([...]) viszont elemeit bejegyzi, ezért a paraméterek megjelennek. A lista és a ModuleList tehát nem felszíni stilisztikai különbség: az egyik a tanulásban részt vevő paraméterek regisztrációját veszíti el.

Önellenőrzés. Ha egy bufferre m.register_buffer('q', t) hívást hajtunk végre, a q tenzor requires_grad értéke hamis; a .grad mezője a backward után is None. A state_dict() csak értékeket, nem autograd gráfot ment; a betöltés levéltensorokat ad vissza. A .to(device) a paramétereket és a buffereket is áthelyezi, de az almodulok közötti névstruktúra változatlan marad. A korábbi 11.1. Tensor és autograd-életciklus fogalmait visszahívva: a Parameter levéltensor, és a .grad felhalmozódás nullázása továbbra is az optimizer feladata.

Korábbi kapcsolat. A 10.2. Affín réteg és aktiváció backward egységben levezetett $Z=XW^\top+b$ mellett a súlygradiens $dW=G^\top X$, a biasgradiens $db=\sum_n G_{n,:}$. Az nn.Linear ugyanezt a transzponált konvenciót használja, ezért a kézi VJP-k változatlanul alkalmazhatók; a modul csupán a tenzorok regisztrációját és nevesítését adja hozzá.

11.2.8 Rövid összefoglaló

Az nn.Module a tenzorokból nevesített, menthető modellállapotot épít. A Parameter a modul _parameters nyilvántartásába kerül, requires_grad=True, ezért az optimizer és a state_dict látja. A register_buffer tenzora a _buffers-be kerül: nem tanul, gradienst nem kap, de a .to() áthelyezi és a mentés tartalmazza. A forward a számítást adja, a train()/eval() a training flaget rekurzívan állítja; a Dropout csak tanításban dobozol. A labor 3-4-2 MLP-n ellenőrizte a 26 paramétert, a kulcsokat, az eszközt és az eval-módú identitást. A Python-listás almodulok kihagyása tipikus regisztrációs hiba.

11.2.9 Következő alkalmazás és források

A következő 11.3. Dataset és DataLoader a split után egységben a betöltőt építjük fel: a modell paraméterei már regisztrálva vannak, most az adat batch-ekre bontása és a split utáni iteráció következik, majd a 11.4. Egy teljes PyTorch tanítási iteráció alfejezetben a teljes tanítási lépés összeáll.

• [D03] PyTorch contributors (2.14). Tensor. https://docs.pytorch.org/docs/2.14/tensors.html.
• [D05] PyTorch contributors (2.14). Autograd mechanics. https://docs.pytorch.org/docs/2.14/notes/autograd.html.