File size: 1,488 Bytes
c68a2e6 | 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 | import torch
import torch.nn as nn
class DoubleConv(nn.Module):
def __init__(self, cin, cout):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(cin, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
nn.Conv2d(cout, cout, 3, padding=1), nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
)
def forward(self, x): return self.net(x)
class UNet(nn.Module):
def __init__(self, in_ch=3, out_ch=3, base=48):
super().__init__()
self.d1 = DoubleConv(in_ch, base)
self.d2 = DoubleConv(base, base * 2)
self.d3 = DoubleConv(base * 2, base * 4)
self.pool = nn.MaxPool2d(2)
self.bottleneck = DoubleConv(base * 4, base * 8)
self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
self.u3 = DoubleConv(base * 8, base * 4)
self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
self.u2 = DoubleConv(base * 4, base * 2)
self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
self.u1 = DoubleConv(base * 2, base)
self.out = nn.Conv2d(base, out_ch, 1)
def forward(self, x):
c1 = self.d1(x)
c2 = self.d2(self.pool(c1))
c3 = self.d3(self.pool(c2))
b = self.bottleneck(self.pool(c3))
x = self.u3(torch.cat([self.up3(b), c3], 1))
x = self.u2(torch.cat([self.up2(x), c2], 1))
x = self.u1(torch.cat([self.up1(x), c1], 1))
return self.out(x)
|