File size: 3,610 Bytes
b8fae22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Attention U-Net (Oktay et al., 2018) — additive attention gates on skips.
Reimplemented cleanly inside the framework. The original sfczekalski/attention_unet
repo is unmaintained (2020), unlicensed, has DRIVE-specific hardcoded padding and a
commented-out training loop, so only the architecture idea is reused here. This is
the faithful additive gating attention (NOT SMP's scSE), configurable for arbitrary
in_channels / num_classes.
"""
from __future__ import annotations
import torch
import torch.nn as nn
class ConvBlock(nn.Module):
def __init__(self, cin, cout):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(cin, cout, 3, padding=1, bias=False),
nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
nn.Conv2d(cout, cout, 3, padding=1, bias=False),
nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
)
def forward(self, x):
return self.block(x)
class UpConv(nn.Module):
def __init__(self, cin, cout):
super().__init__()
self.up = nn.Sequential(
nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
nn.Conv2d(cin, cout, 3, padding=1, bias=False),
nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
)
def forward(self, x):
return self.up(x)
class AttentionGate(nn.Module):
"""Additive attention gate: gating signal g (coarser) modulates skip x."""
def __init__(self, f_g, f_l, f_int):
super().__init__()
self.w_g = nn.Sequential(nn.Conv2d(f_g, f_int, 1, bias=True), nn.BatchNorm2d(f_int))
self.w_x = nn.Sequential(nn.Conv2d(f_l, f_int, 1, bias=True), nn.BatchNorm2d(f_int))
self.psi = nn.Sequential(nn.Conv2d(f_int, 1, 1, bias=True), nn.BatchNorm2d(1), nn.Sigmoid())
self.relu = nn.ReLU(inplace=True)
def forward(self, g, x):
a = self.relu(self.w_g(g) + self.w_x(x))
a = self.psi(a)
return x * a
class AttentionUNet(nn.Module):
def __init__(self, in_channels: int = 3, num_classes: int = 2, base: int = 64):
super().__init__()
c = [base, base * 2, base * 4, base * 8, base * 16]
self.pool = nn.MaxPool2d(2, 2)
self.e1 = ConvBlock(in_channels, c[0])
self.e2 = ConvBlock(c[0], c[1])
self.e3 = ConvBlock(c[1], c[2])
self.e4 = ConvBlock(c[2], c[3])
self.e5 = ConvBlock(c[3], c[4])
self.u5 = UpConv(c[4], c[3]); self.a5 = AttentionGate(c[3], c[3], c[2]); self.d5 = ConvBlock(c[4], c[3])
self.u4 = UpConv(c[3], c[2]); self.a4 = AttentionGate(c[2], c[2], c[1]); self.d4 = ConvBlock(c[3], c[2])
self.u3 = UpConv(c[2], c[1]); self.a3 = AttentionGate(c[1], c[1], c[0]); self.d3 = ConvBlock(c[2], c[1])
self.u2 = UpConv(c[1], c[0]); self.a2 = AttentionGate(c[0], c[0], c[0] // 2); self.d2 = ConvBlock(c[1], c[0])
self.head = nn.Conv2d(c[0], num_classes, 1)
def forward(self, x):
x1 = self.e1(x)
x2 = self.e2(self.pool(x1))
x3 = self.e3(self.pool(x2))
x4 = self.e4(self.pool(x3))
x5 = self.e5(self.pool(x4))
d5 = self.u5(x5); d5 = self.d5(torch.cat([self.a5(d5, x4), d5], dim=1))
d4 = self.u4(d5); d4 = self.d4(torch.cat([self.a4(d4, x3), d4], dim=1))
d3 = self.u3(d4); d3 = self.d3(torch.cat([self.a3(d3, x2), d3], dim=1))
d2 = self.u2(d3); d2 = self.d2(torch.cat([self.a2(d2, x1), d2], dim=1))
return self.head(d2)
def build_attention_unet(in_channels: int, num_classes: int, **_):
return AttentionUNet(in_channels=in_channels, num_classes=num_classes)
|