Segmentation Architecture: U-Net

နားလည်ထားရမယ့် အချက်

Classification network တွေဟာ deeper သွားလေလေ spatial resolution ကို တမင်ချန်ထားခဲ့ပါတယ်— pooling နဲ့ strided convolution တွေက feature map ကို သေးငယ်၊ abstract ကောင်းစွာရှိတဲ့ representation တစ်ခုအဖြစ် ချုံ့ပေးပြီး၊ နောက်ဆုံး layer တစ်ခုက ပုံတစ်ပုံစီအတွက် label တစ်ခုတည်းအဖြစ် ချုံ့ချလိုက်ပါတယ်။ Segmentation က ဆန့်ကျင်ဘက် output လိုချင်ပါတယ်— pixel တစ်ခုစီအတွက် label တစ်ခု၊ original image resolution အတိုင်းလိုအပ်ပါတယ်။ encoder-only network ရိုးရိုးတစ်ခုက ဒါကို ကောင်းကောင်း မလုပ်နိုင်ပါဘူး၊ feature တွေက 'ဒီနေရာက ကြောင်တစ်ကောင်' ဆိုတာသိလောက်အောင် abstract ဖြစ်လာချိန်မှာ ကြောင်ရဲ့ edge တွေ အတိအကျ ဘယ်နေရာလဲဆိုတဲ့ fine spatial detail အားလုံး pooling ကြောင့် ပျောက်ဆုံးပြီးသားဖြစ်နေလို့ပါ။ U-Net ရဲ့ architecture ဟာ symmetric encoder-decoder တစ်ခုပါ— encoder (contracting path) က ပုံမှန်အတိုင်း downsample လုပ်ပြီး semantic understanding တည်ဆောက်ပြီး၊ mirror ဖြစ်တဲ့ decoder (expansive path) က transposed convolution ဒါမှမဟုတ် interpolation သုံးပြီး original resolution ဆီ တစ်ဆင့်ချင်း upsample ပြန်လုပ်ပါတယ်။

U-Net ကို blur ဖြစ်နေတဲ့ upsampled blob တစ်ခုတည်း ထုတ်ပေးရုံမက တကယ် ကောင်းကောင်း အလုပ်လုပ်စေတဲ့ အစိတ်အပိုင်းကတော့ mirror ဖြစ်တဲ့ encoder နဲ့ decoder stage တွေကြားက skip connection တွေပါပဲ— decoder step တစ်ခုစီမှာ upsample လုပ်ထားတဲ့ feature map ကို convolution ထပ်ဖြတ်မခင် (ResNet လိုမျိုး addition မဟုတ်ဘဲ) spatial resolution တူညီတဲ့ encoder feature map နဲ့ concatenate လုပ်ပါတယ်။ ဒါက decoder ကို encoder ထဲမှာ အစောပိုင်းက ဖမ်းယူထားတဲ့ high-resolution spatial detail — အစွန်းအနား sharp တွေနဲ့ fine texture တွေ — ကို တိုက်ရိုက် access ပေးထားပြီး bottleneck ကနေ လာတဲ့ coarse ဒါပေမယ့် semantic ကြွယ်ဝတဲ့ feature တွေကလည်း အကျိုးရရှိစေပါတယ်။ ဒီ skip connection တွေ မပါရင် decoder ဟာ heavily compress လုပ်ထားတဲ့ bottleneck representation တစ်ခုတည်းကနေ pixel boundary အတိအကျကို ပြန်တည်ဆောက်ရမှာဖြစ်ပြီး၊ ဒါဟာ spatial information အများကြီး ဆုံးရှုံးထားလို့ တိကျစွာ လုပ်ဖို့ မဖြစ်နိုင်ပါဘူး။

လက်တွေ့ scenario နဲ့ ချိတ်ကြည့်မယ်

Tutorial Platform မှာ user တွေ submit လုပ်လိုက်တဲ့ code screenshot များထဲက sensitive region များ — API key, token, ဒါမှမဟုတ် terminal ထဲကို paste လုပ်လိုက်တဲ့ personal file path များ — ကို pixel-precise mask တစ်ခု ထုတ်ပေးဖို့ U-Net-style segmentation model ကို သုံးနိုင်ပါတယ်၊ full-resolution pixel-level output (coarse bounding box တစ်ခုတည်း မဟုတ်ဘဲ) ရှိမှသာ screenshot ကို publish မလုပ်ခင် ပတ်ဝန်းကျင်က code ကို မဖုံးအုပ်မိစေဘဲ ပြဿနာရှိတဲ့ text အတိအကျကိုပဲ blur လုပ်ပေးနိုင်မှာမို့ပါ။

Classification network တွေဟာ deeper သွားလေလေ spatial resolution ကို စွန့်ပစ်လေလေ ဖြစ်လို့ pixel-level segmentation output အတွက် plain encoder-only network တစ်ခုတည်း မသင့်တော်ပါဘူး။
U-Net ရဲ့ decoder ဟာ encoder ကို mirror လုပ်ပြီး label တစ်ခုတည်း ချုံ့ချမယ့်အစား original spatial resolution ဆီ တစ်ဆင့်ချင်း upsample ပြန်လုပ်ပါတယ်။
Skip connection တွေဟာ resolution တူညီတဲ့ encoder feature တွေကို decoder ထဲကို concatenate (addition မဟုတ်ဘဲ) လုပ်ပေးပြီး downsampling အတွင်း ဆုံးရှုံးသွားတဲ့ fine spatial detail ကို ပြန်ရရှိစေပါတယ်။

အတူတူ စမ်းရေးကြည့်မယ်

python

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

class MiniUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2d(2)
        self.enc2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2d(2)

        self.bottleneck = nn.Conv2d(32, 64, kernel_size=3, padding=1)

        self.up2 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.dec2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)   # 32 (up2) + 32 (skip) = 64 in
        self.up1 = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
        self.dec1 = nn.Conv2d(32, 16, kernel_size=3, padding=1)   # 16 (up1) + 16 (skip) = 32 in

        self.out_conv = nn.Conv2d(16, 1, kernel_size=1)

    def forward(self, x):
        e1 = F.relu(self.enc1(x))          # (batch, 16, H, W)
        p1 = self.pool1(e1)                # (batch, 16, H/2, W/2)
        e2 = F.relu(self.enc2(p1))         # (batch, 32, H/2, W/2)
        p2 = self.pool2(e2)                # (batch, 32, H/4, W/4)

        b = F.relu(self.bottleneck(p2))    # (batch, 64, H/4, W/4)

        u2 = self.up2(b)                                       # (batch, 32, H/2, W/2)
        d2 = F.relu(self.dec2(torch.cat([u2, e2], dim=1)))      # skip connection with e2

        u1 = self.up1(d2)                                      # (batch, 16, H, W)
        d1 = F.relu(self.dec1(torch.cat([u1, e1], dim=1)))      # skip connection with e1

        return self.out_conv(d1)           # (batch, 1, H, W)

torch.manual_seed(0)
model = MiniUNet()
images = torch.randn(2, 3, 64, 64)  # fake batch of 2 RGB images
output = model(images)
print(output.shape)
print(output.shape[-2:] == images.shape[-2:])

You should see

torch.Size([2, 1, 64, 64])
True — downsampling stage နှစ်ခု (64x64 -> 32x32 -> 16x16) နဲ့ ကိုက်ညီတဲ့ upsampling stage နှစ်ခု (16x16 -> 32x32 -> 64x64) ပြီးနောက် output spatial dimension ဟာ 64x64 input နဲ့ အတိအကျ ကိုက်ညီပြီး၊ နောက်ဆုံး print က ဒီညီမျှမှုကို အတည်ပြုပါတယ်— output channel တစ်ခုတည်းက pixel-level prediction map တစ်ခုကို ကိုယ်စားပြုပြီး၊ ဒါက U-Net ကို segmentation အတွက် သင့်တော်စေတဲ့ အဓိက property ပါပဲ။

၅ မိနစ် စမ်းကြည့်

encoder/decoder level တတိယတစ်ခု (H/8 ကို downsample လုပ်မယ့် enc3/pool3၊ e3 ကို skip connection ချိတ်ထားတဲ့ ကိုက်ညီတဲ့ up3/dec3) ကို ထပ်ထည့်ပြီး နောက်ဆုံး print ထုတ်တဲ့ output shape ဟာ (2, 1, 64, 64) အတိုင်း ကျန်ရှိနေသေးလားစစ်ဆေးပါ— ဒါက downsampling stage တစ်ခုစီမှာ mirror ဖြစ်တဲ့ upsampling stage ရှိနေသရွေ့ skip-connection pattern ဟာ U-Net မည်မျှပင် နက်ရှိုင်းနေသော် scale ဖြစ်နိုင်ကြောင်း ပြသပါတယ်။

သတိလေးတစ်ချက်

pool layer နှစ်ခုကနေ downsampling factor စုစုပေါင်း 4 နဲ့ အညီအမျှ မစားနိုင်တဲ့ input spatial size သုံးလိုက်ရင် encoder နဲ့ decoder feature map တွေရဲ့ size တွေ မကိုက်ညီတော့ဘဲ skip connection မှာ torch.cat က dimension-mismatch error တက်ပါလိမ့်မယ်။

skip connection မှာ torch.cat အစား ရိုးရှင်းတဲ့ addition ကို (ResNet ကို အတုယူပြီး) အစားထိုးလိုက်ရင် error မတက်ဘဲ တိတ်တဆိတ် run သွားနိုင်ပေမယ့် information ဆုံးရှုံးသွားပါလိမ့်မယ်၊ encoder နဲ့ decoder feature set နှစ်ခုကို အနားချင်းယှဉ် ထိန်းသိမ်းမယ့်အစား channel count တွေ ညီမျှနေရအောင် အတင်း လုပ်စေလို့ပါ။

Wikipedia — U-Net — Computer Vision