27.2 U-Net, encoder–decoder és skip

27.2.1 1. Cél, előismeret és induló kérdés

A 27.1. Pixelszintű klasszifikáció és FCN egységben a szegmentálást pixelszintű osztályozásként definiáltuk: a modell minden pixelre logitot ad, a tanítás per-pixel cross-entropy vagy bináris logitveszteséggel történik. Ott a kérdés az volt, hogyan kell a logitok és a maszk tengelyeit párosítani. Most azt vizsgáljuk, milyen hálózati szerkezet állítja elő a térbeli kimenetet anélkül, hogy a felbontást azonnal eldobnánk. A U-Net az egyik leggyakrabban használt encoder–decoder architektúra; benne a downsampling kontextust gyűjt, az upsampling a felbontást állítja helyre, a skip-kapcsolat pedig a korai, nagy felbontású aktivációkat közvetlenül a dekódolóhoz vezeti.

Előismeret: a 15.2. Inputgradiens és átfedő recepciós mezők alfejezetből a konvolúciós inputgradiens scatter-add indexfeltétele, a 15.3. Transposed convolution mint adjungált alfejezetből a transzponált konvolúció mint adjungált operátor értelmezése, a 23.2. Kalibráció és döntési küszöb alfejezetből a logit és döntési küszöb különválasztása, a 26.2. Egy- és kétlépcsős detektor alfejezetből pedig a detektálási fejek felelősségének szétválasztása. A U-Net itt nem javaslattevő rendszer: maszkot jósol.

Induló kérdés. Egy bemeneti kép $24\times24$-es, egy csatornás. A kódoló két $2\times2$-es max-pool-lépéssel $6\times6$-ra csökkenti a felbontást. A dekódoló két tanulható upsampling-lépést használ. Mekkora lesz a kimenet, és hány csatornája van a második konkatenációnál, ha a skip-ek rendre $8$ és $16$ csatornásak? Becslés után a 3. szakasz levezetése ad pontos választ.

27.2.2 2. Definíciók, jelölések és tenzoralakok

Legyen a bemeneti képbatch $X\in\mathbb R^{N\times C_{in}\times H\times W}$, ahol a főútban $C_{in}=1$, $H=W=24$. A bináris szegmentálási maszk $Y\in\{0,1\}^{N\times 1\times H\times W}$: minden pixelhez egyetlen logit tartozik. A hálózat kimenete $Z\in\mathbb R^{N\times 1\times H\times W}$, ahol $Z[n,0,i,j]$ a $(n,i,j)$ pixel előtér-logitja, szigmoid előtt.

Encoder alatt azt a részhálózatot értjük, amely a térbeli felbontást csökkenti, miközben a csatornaszámot növeli. Decoder az a részhálózat, amely a felbontást ismét növeli és a csatornaszámot csökkenti. Skip concat a kódoló egy közbülső aktivációjának és a dekódoló azonos térbeli felbontású aktivációjának csatorna-irányú összefűzése:

$$Z_{cat} = \operatorname{cat}\bigl(U, S\bigr) \quad\text{a } C\text{ dimenzió mentén}, \qquad Z_{cat}\in\mathbb R^{N\times(C_U+C_S)\times H'\times W'}.$$

A torch.cat hívásban ez dim=1, mert a tenzoralak NCHW.

Upsampling minden olyan leképezés, amely $H,W$ térbeli méretet növel. Két fő változatot különböztetünk meg. A nem tanulható interpoláció, például F.interpolate(..., mode="bilinear"), rögzített szabállyal becsli a közbülső értékeket; ennek nincs paramétere. A tanulható ConvTranspose2d a konvolúciós operátor adjungáltjaként szór szét aktivációkat; ez nem általános inverz, hanem lineáris leképezés, amely a 15.3. Transposed convolution mint adjungált alfejezetben tárgyalt alakkal rendelkezik. A jelen fejezetben $k=2$, $s=2$, $p=0$ transzponált konvolúciót használunk, amely minden $H$-t $2H$-ra növel.

Spatial alignment az a feltétel, hogy a skip tenzor $H_S\times W_S$ kiterjedése egyenlő legyen az upsampling kimenet $H_U\times W_U$ kiterjedésével a cat előtt. Ha nem egyenlő, akkor crop vagy pad szükséges; a mi főpéldánkban a $24$, $4$-gyel osztható bemenet miatt a feltétel teljesül.

27.2.3 3. Alaklevezetés egy kis U-Netre

A levezetési kötelezettség egy konkrét, teljesen megadott kis U-Net minden térbeli alakjának és a két concat csatornaszámának kiszámítása. A modellt TinyUNet-nek nevezzük; ez a könyv saját tanulópéldája, nem az eredeti publikáció másolata. Térbeli konvolúciók: $k=3$, padding $p=1$, stride $s=1$, dilation $d=1$. Pool: $2\times2$ max, stride $2$. Upsampling: $k=2$, $s=2$, $p=0$ transposed conv.

Lemma (alak- és csatornaegyezés a TinyUNetben). Legyen $X\in\mathbb R^{N\times1\times24\times24}$. A fenti paraméterezésű TinyUNet esetén

1. a kódoló kimenetei rendre $24\times24$, $12\times12$, $6\times6$ térbeli méretűek;
2. az első upsampling kimenete $12\times12$, a skip hozzáadása után a csatornaszám $16+16=32$;
3. a második upsampling kimenete $24\times24$, a skip hozzáadása után a csatornaszám $8+8=16$;
4. ezért egyik concat előtt sem kell térbeli crop vagy pad.

Bizonyítás. A konvolúciós kimeneti méret képlete 14.1. Stride, padding, dilation és outputméret jelölésével, $d=1$ mellett:

$$H_{out}=\left\lfloor\frac{H+2p-(k-1)-1}{s}\right\rfloor+1.$$

A $k=3$, $p=1$, $s=1$ esetben

$$H_{out}=\left\lfloor\frac{H+2-2-1}{1}\right\rfloor+1 =\left\lfloor H-1\right\rfloor+1=H,$$

tehát az ilyen konvolúció a térbeli méretet megtartja. A $2\times2$-es, stride-2 max-pool képlete $p=0$, $k=2$ esetén

$$H_{out}=\left\lfloor\frac{H-2}{2}\right\rfloor+1.$$

Innen $H=24$ esetén $H_{out}=12$, majd $H=12$ esetén $H_{out}=6$. A transzponált konvolúció $p=0$, $d=1$, $k=2$, $s=2$ paraméterekkel a 15.3. Transposed convolution mint adjungált alfejezetben levezetett képlet szerint

$$H_{out}=(H-1)s-2p+d(k-1)+\text{output\_padding}+1 =(H-1)\cdot2+1+0+1=2H.$$

Ezért a $6\times6$ bottleneck-ből $12\times12$, majd a $12\times12$-ből $24\times24$ lesz. A cat a csatornadimenzióban ad össze: az első upsampling $16$, a hozzá tartozó skip $16$ csatornás, így $32$; a második upsampling $8$, a skip $8$ csatornás, így $16$. A térbeli méretek a fenti értékek szerint rendre megegyeznek, tehát az alignment feltétel teljesül. ∎

A teljes rétegenkénti alak a következő táblázatban:

A táblázat egyben kézi példa is: a $12\times12$-es skip és a $12\times12$-es up-aktiváció csatorna szerint $32$-re bővül, majd a $24\times24$-es szinten $16$-ra. A bottlenecknél a térbeli felbontás $6\times6$, a csatorna pedig $32$: ez a modell legsűrűbb csatornareprezentációja.

27.2.4 4. Eredeti kézi példa és minimális kód

A levezetést a következő, teljesen futtatható PyTorch-modullal ellenőrizhetjük. A modul minden Conv2d rétegnél padding=1-et használ, a két ConvTranspose2d pedig kernel_size=2, stride=2 paraméterű, ezért a fenti lemmát valósítja meg.

# tiny_unet_shapecheck.py
import torch
import torch.nn as nn

class TinyUNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, base=8):
        super().__init__()
        # encoder, 1. szint: 1 -> 8 csatorna, 24x24 megmarad
        self.enc1a = nn.Conv2d(in_channels, base, 3, padding=1)
        self.enc1b = nn.Conv2d(base, base, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2)
        # encoder, 2. szint: 8 -> 16, 12x12 megmarad
        self.enc2a = nn.Conv2d(base, 2 * base, 3, padding=1)
        self.enc2b = nn.Conv2d(2 * base, 2 * base, 3, padding=1)
        self.pool2 = nn.MaxPool2d(2)
        # bottleneck: 16 -> 32, 6x6
        self.bot_a = nn.Conv2d(2 * base, 4 * base, 3, padding=1)
        self.bot_b = nn.Conv2d(4 * base, 4 * base, 3, padding=1)
        # decoder, 1. szint
        self.up1 = nn.ConvTranspose2d(4 * base, 2 * base, 2, stride=2)
        self.dec1a = nn.Conv2d(4 * base, 2 * base, 3, padding=1)  # 16 skip + 16 up
        self.dec1b = nn.Conv2d(2 * base, 2 * base, 3, padding=1)
        # decoder, 2. szint
        self.up2 = nn.ConvTranspose2d(2 * base, base, 2, stride=2)
        self.dec2a = nn.Conv2d(2 * base, base, 3, padding=1)      # 8 skip + 8 up
        self.dec2b = nn.Conv2d(base, base, 3, padding=1)
        # maszkfej: logit, nem softmax előtti valószínűség
        self.head = nn.Conv2d(base, out_channels, 1)

    def forward(self, x):
        e1 = torch.relu(self.enc1b(torch.relu(self.enc1a(x))))       # 8x24x24
        e2 = torch.relu(self.enc2b(torch.relu(self.enc2a(self.pool1(e1)))))  # 16x12x12
        b  = torch.relu(self.bot_b(torch.relu(self.bot_a(self.pool2(e2)))))  # 32x6x6
        u1 = self.up1(b)                                             # 16x12x12
        d1 = torch.relu(self.dec1b(torch.relu(self.dec1a(torch.cat([u1, e2], dim=1)))))
        u2 = self.up2(d1)                                            # 8x24x24
        d2 = torch.relu(self.dec2b(torch.relu(self.dec2a(torch.cat([u2, e1], dim=1)))))
        return self.head(d2)                                         # 1x24x24 logit

A modell célja a rétegspecifikáció ellenőrzése, nem a tanítás. A paraméterszám is kiszámítható a rétegenkénti $(C_{in}C_{out}k_hk_w+C_{out})$ összegből; a fenti modulra ez $29\,321$. Ez magában foglalja a bias-eket is. A dec1a bemeneti csatornája azért $32$, mert a concat után $16+16$ csatorna érkezik; a dec2a bemenete pedig $8+8=16$.

27.2.5 5. Labor: futtatható rétegspecifikáció és várt kimenet

A labor célja annak igazolása, hogy a lemmában szereplő minden térbeli alak és concat-csatornaszám ténylegesen megvalósul a kódban. A következő ellenőrzőkód közvetlenül a fenti modul után illeszthető:

if __name__ == "__main__":
    torch.manual_seed(20261020)
    device = torch.device("cpu")
    model = TinyUNet(in_channels=1, out_channels=1, base=8).to(device)
    model = model.to(torch.float32)

    x = torch.randn(2, 1, 24, 24, dtype=torch.float32, device=device)
    with torch.no_grad():
        model.eval()
        y = model(x)

    print("input", tuple(x.shape))
    print("output", tuple(y.shape))
    print("parameters", sum(p.numel() for p in model.parameters()))

    assert tuple(y.shape) == (2, 1, 24, 24)
    assert sum(p.numel() for p in model.parameters()) == 29321

Futtatás a projektkönyvtárból:

python tiny_unet_shapecheck.py

Várt, számított kimenet — ez nem mért futás, hanem a fenti képletekből és a paraméterszám-összegzésből származó ellenőrzött érték; a kódot a hallgató a saját környezetében futtatja:

input (2, 1, 24, 24)
output (2, 1, 24, 24)
parameters 29321

Az eredmény értelmezése: a kimenet $N\times1\times24\times24$, tehát a hálózat valóban dense predictiont ad: minden bemeneti pixelhez egy logit tartozik. A model.eval() hívás itt azért szerepel, hogy a kiértékelési mód szokását rögzítsük; a modulban nincs BatchNorm vagy dropout, így a train/eval viselkedés most numerikusan nem tér el. A tanításhoz a maszkot float32 típusra konvertálva torch.nn.functional.binary_cross_entropy_with_logits(logits, mask_float) hívható; a logitokból a döntési maszk $\hat Y=\mathbf 1[\sigma(Z)>0{,}5]$ szabállyal áll elő, a 23.2. Kalibráció és döntési küszöb alfejezetben tárgyalt küszöb–kalibráció szétválasztás jegyében.

27.2.6 6. Feladatok

F1 (fogalmi, nehézség 1). Adja meg két konkrét különbséget a bilineáris F.interpolate-alapú upsampling és a ConvTranspose2d-alapú upsampling között a tanulhatóság és az adjungált jelleg szempontjából. Előfeltétel: 15.3. Transposed convolution mint adjungált. Átadandó: legfeljebb fél oldalas saját válasz, amely mindkét állítást a jelen modul példájára vonatkoztatja.

F2 (levezetési, nehézség 2). Vezesse le a fenti TinyUNet rétegalakjait $H=W=22$ bemenetre. Hol sérül a spatial alignment, és mekkora a második upsampling kimenete? Előfeltétel: a 3. szakasz formulái. Átadandó: a pool- és upsampling-méretek lépésenként, a szükséges crop mérete, valamint annak kimondása, hogy a kimenet most nem $22\times22$.

F3 (implementációs, nehézség 1). Írjon center_crop(x, target_h, target_w) függvényt, amely az utolsó két dimenziót a megadott méretre vágja középre, és tesztelje $x\in\mathbb R^{2\times8\times23\times23}$, target_h=target_w=22 bemenetre. Átadandó: kód, futtatási parancs és az ellenőrzött kimeneti alak.

F4 (hibakeresési, nehézség 2). A következő modulrészlet az első decoder szint definíciója:

self.up1 = nn.ConvTranspose2d(32, 16, 2, stride=2)
self.dec1a = nn.Conv2d(16, 16, 3, padding=1)
# ...
d1 = torch.cat([self.up1(b), e2], dim=1)
d1 = torch.relu(self.dec1a(d1))

Mi a hiba, mi a közvetlen következménye, és hogyan javítható? Átadandó: ok, következmény, a javított sor és a visszaellenőrzés módja.

F5 (kísérleti/transzfer, nehézség 3). Tervezzen GlyphLab-kísérletet, amely a skip-kapcsolatok hatását vizsgálja a szegmentálásra. A cél: ugyanaz a TinyUNet skip-pel és skip nélkül. A maszk a data/masks.npy előtér-maszkja. Adjon meg: megfigyelési egység, split, baseline, változó tényező, loss, metrika, a választás szabálya és a következtetés határa. Előfeltétel: 27.1. Pixelszintű klasszifikáció és FCN, 1.3. Adatfelosztás és a megfigyelési egység. Átadandó: protokoll, nem várt győztes.

27.2.7 7. Megoldások és indoklás

F1 megoldás. A bilineáris interpoláció paraméter nélküli, determinisztikus leképezés: a kimeneti érték a bemeneti értékek rögzített súlyozott átlaga, ezért a hálózat a tanítás során nem tudja a feladathoz igazítani. A ConvTranspose2d ezzel szemben tanulható kernelű lineáris leképezés. Második különbség: a ConvTranspose2d a megfelelő konvolúciós lineáris operátor adjungáltja, tehát a forward képlete a 15.3. Transposed convolution mint adjungált alfejezetben levezetett scatter-add szerkezetű; a bilineáris interpoláció nem ilyen adjungált szerkezetű. A jelen példában a up1 $32\to16$ csatornás tanulható leképezés, miközben a bilineáris változat legfeljebb a meglévő csatornákat interpolálná, csatornát nem keverne.

F2 megoldás. $H=22$ esetén a kódolási lépések:

$$22 \xrightarrow{\text{conv }3\times3,\ p=1} 22 \xrightarrow{\text{pool}} \left\lfloor\frac{22-2}{2}\right\rfloor+1=11 \xrightarrow{\text{pool}} \left\lfloor\frac{11-2}{2}\right\rfloor+1=5.$$

Az első upsampling $5$-ből

$$(5-1)\cdot 2+1+1=10$$

méretet állít elő, a skip viszont $11\times11$, tehát $1$ pixel eltérés van: a skip-ből $10\times10$-re kell vágni. A concat csatornaszáma változatlanul $16+16=32$. A második upsampling $10$-ből $20$-at állít elő, a skip itt $22\times22$, tehát $2$ pixelnyi eltérés van. A kimenet ezért $20\times20$, nem $22\times22$. A pontos crop-policy implementációfüggő; a levezetés állítása csak a méreteltérés és a kimeneti felbontás. Ez mutatja a lemma feltételének szerepét: a $24$-es bemenet azért igazodik, mert $24/4=6$ egész.

F3 megoldás.

def center_crop(x, target_h, target_w):
    h, w = x.shape[-2:]
    top = (h - target_h) // 2
    left = (w - target_w) // 2
    return x[..., top:top + target_h, left:left + target_w]

# teszt
import torch
x = torch.zeros(2, 8, 23, 23)
y = center_crop(x, 22, 22)
assert tuple(y.shape) == (2, 8, 22, 22)

A függvény a top, left levágással a felesleget a bal/felső és jobb/alsó oldal között osztja el; $23\to22$-nél (23-22)//2=0, tehát a jobb/alsó egy pixelt elveszít. A függvény a feladat specifikációjának megfelel; ha szigorúan szimmetrikus kerekítés kell, azt a feladatban külön definiálni kell.

F4 megoldás. A hiba oka: a cat([up1(b), e2], dim=1) kimenete $32$ csatornás, mert az upsampling $16$, a skip $16$ csatornát ad össze. A dec1a azonban $16$ bemeneti csatornát deklarál. Következmény: a PyTorch a Conv2d bemeneti csatorna-ellenőrzésénél RuntimeError-t ad, mert a várt $16$ helyett $32$ csatorna érkezik. Javítás: self.dec1a = nn.Conv2d(32, 16, 3, padding=1), vagy 4 * base jelöléssel self.dec1a = nn.Conv2d(4 * base, 2 * base, 3, padding=1). Visszaellenőrzés: a 4. szakasz kódjának assert-jein a kimenet $(2,1,24,24)$ és a paraméterszám $29\,321$ marad, mert a javított rétegnek $32\cdot16\cdot9+16=4624$ paramétere van.

F5 megoldásminta. Kérdés: a skip-kapcsolatok javítják-e a GlyphLab előtér-szegmentálást? Megfigyelési egység: scene_id, mert egy jelenet két nézete függő; a train/val/test split a GlyphLab szerződés szerint jelenetszintű. Bemenet: data/images.npy float32 $(360,1,24,24)$; cél: data/masks.npy uint8 $(360,1,24,24)$. Baseline: ugyanaz a TinyUNet, amelyben a két cat-et eltávolítjuk — a skip nélküli decoder csak a bottleneckből építkezik. Változó tényező: skip megléte. Loss: binary_cross_entropy_with_logits, a maszk float32-vé konvertálva. Metrika: threshold $0{,}5$ melletti pixel-IoU képenként, majd a jelenetszintű aggregálás a 27.3. Dice, IoU és aggregálás alfejezetben tárgyalandó módon; a képenkénti IoU-kat itt még nem szabad független mintaként kezelni. Választási szabály: a legjobb kísérletenkénti checkpointot a validációs loss választja, a test-halmazt csak a végleges modellpár kiválasztása után értékeljük. Erőforrás: CPU, float32, rögzített seed. Következtetési határ: a GlyphLab szintetikus, három orientációhoz tartozó alakzatmaszk; az eredmény nem általános természetes képi szegmentálási validáció, és a skip hatása csak ezen az adaton, ezzel a protokollal állítható.

27.2.8 8. Tipikus tévedés, önellenőrzés és visszahívás

A gyakori tévedés az, hogy a fenti TinyUNet-et az eredeti U-Net reprodukciójának nevezzük. A jelen modul U-Net-szerű tanulópélda: van encoder, decoder és skip concat, de az eredeti, [P14] által hivatkozott U-Net orvosbiológiai képszegmentálásra készült, eltérő bemeneti felbontással, mélységgel, csatornaszámmal és tanítási eljárással. Az a kijelentés, hogy „implementáltuk a U-Net cikket”, csak akkor lenne igazolható, ha a modell, az adat, a loss, az augmentáció és a kiértékelési metrika a hivatkozott közleményben leírt protokollnak felelne meg; ezek jelen kéziratban nincsenek reprodukálva. A mi ellenpéldánk a konkrét adatszerződés: a GlyphLab $360$ darab $24\times24$-es szintetikus kép, míg a [P14] absztraktja orvosbiológiai szegmentálási felhasználásról szól. A két állítás közötti különbség nem pusztán méretbeli, hanem a reprodukciós állítás típusában van.

Önellenőrzés: minden $3\times3$-as, $p=1$ konvolúció megtartja a térbeli méretet? A kódban minden decoder-Conv2d is padding=1? A két cat valóban dim=1-en történik? A dec1a bemeneti csatornája $32$, a dec2a bemenete $16$? Ha a válasz bármelyikre nem, a paraméterszám-ellenőrzés vagy a shape-assert el fog csúszni.

Visszahívás a 15.2. Inputgradiens és átfedő recepciós mezők alfejezetből: a transzponált konvolúció upsamplingjének magja ugyanaz a scatter-add gondolat, mint az inputgradiens számításakor. A $k=2$, $s=2$ esetben egy $(i,j)$ bemeneti pixel a kimeneti $(2i+u,2j+v)$, $u,v\in\{0,1\}$ pozíciókhoz járul hozzá; ezért a kimenet $2H\times2W$ lesz. A skip concat ezt a magasabb felbontású, kódoló oldali aktivációval egészíti ki, de nem változtat a scatter-add indexfeltételén.

27.2.9 Rövid összefoglaló

A U-Net encoder–decoder szerkezet: a kódoló konvolúciókkal és max-poollal növeli a csatornaszámot, miközben csökkenti a térbeli felbontást; a dekódoló tanulható transzponált konvolúciókkal növeli a felbontást. A skip concat a kódoló azonos felbontású aktivációját csatorna szerint fűzi a dekódolóhoz, így a korai térbeli részletek nem csak a bottlenecken át jutnak a kimenetre. A lemmában levezettük a $24\times24$-es bemenet minden rétegalakját: $24\to12\to6\to12\to24$, a concat-csatornaszámok $32$ és $16$. A $22\times22$-es ellenpélda megmutatta, hogy a pool és upsampling méretezése miatt crop válhat szükségessé. A $29\,321$ paraméteres modul a szerződésnek megfelelő rétegspecifikáció.

27.2.10 Következő alkalmazás

A 27.3. Dice, IoU és aggregálás egységben a szegmentálás értékelését építjük fel: a Dice-együttható és az IoU definícióját, a per-image, per-scene és micro aggregálás különbségét, valamint az üres maszkok és a küszöbválasztás hatását. A mostani Tensor-logit maszk az ottani metrikák közvetlen bemenete.

27.2.11 Felhasznált források

• [P14] Ronneberger, Fischer és Brox (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. arXiv:1505.04597. https://arxiv.org/abs/1505.04597.
• [P15] He et al. (2017). Mask R-CNN. arXiv:1703.06870. https://arxiv.org/abs/1703.06870.