File size: 1,836 Bytes
ecc7dcb | 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 | class DoubleConv(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.block(x)
class UNet(nn.Module):
def __init__(self, in_ch=3, out_ch=1, base=64):
super().__init__()
c = [base, base * 2, base * 4, base * 8, base * 16]
self.enc1 = DoubleConv(in_ch, c[0])
self.enc2 = DoubleConv(c[0], c[1])
self.enc3 = DoubleConv(c[1], c[2])
self.enc4 = DoubleConv(c[2], c[3])
self.bottleneck = DoubleConv(c[3], c[4])
self.pool = nn.MaxPool2d(2)
self.up4 = nn.ConvTranspose2d(c[4], c[3], 2, stride=2)
self.dec4 = DoubleConv(c[4], c[3])
self.up3 = nn.ConvTranspose2d(c[3], c[2], 2, stride=2)
self.dec3 = DoubleConv(c[3], c[2])
self.up2 = nn.ConvTranspose2d(c[2], c[1], 2, stride=2)
self.dec2 = DoubleConv(c[2], c[1])
self.up1 = nn.ConvTranspose2d(c[1], c[0], 2, stride=2)
self.dec1 = DoubleConv(c[1], c[0])
self.head = nn.Conv2d(c[0], out_ch, 1)
def forward(self, x):
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
e4 = self.enc4(self.pool(e3))
b = self.bottleneck(self.pool(e4))
d4 = self.dec4(torch.cat([self.up4(b), e4], 1))
d3 = self.dec3(torch.cat([self.up3(d4), e3], 1))
d2 = self.dec2(torch.cat([self.up2(d3), e2], 1))
d1 = self.dec1(torch.cat([self.up1(d2), e1], 1))
return self.head(d1) |