File size: 2,368 Bytes
aa04f76 |
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 |
import torch
from torch import nn
class DoubleConv(nn.Module):
def __init__(self, in_channels, out_channels, with_bn=False):
super().__init__()
if with_bn:
self.step = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
)
else:
self.step = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(),
)
def forward(self, x):
return self.step(x)
class UNet(nn.Module):
def __init__(self, in_channels, out_channels, with_bn=False):
super().__init__()
init_channels = 32
self.out_channels = out_channels
self.en_1 = DoubleConv(in_channels , init_channels , with_bn)
self.en_2 = DoubleConv(1*init_channels, 2*init_channels, with_bn)
self.en_3 = DoubleConv(2*init_channels, 4*init_channels, with_bn)
self.en_4 = DoubleConv(4*init_channels, 8*init_channels, with_bn)
self.de_1 = DoubleConv((4 + 8)*init_channels, 4*init_channels, with_bn)
self.de_2 = DoubleConv((2 + 4)*init_channels, 2*init_channels, with_bn)
self.de_3 = DoubleConv((1 + 2)*init_channels, 1*init_channels, with_bn)
self.de_4 = nn.Conv2d(init_channels, out_channels, 1)
self.maxpool = nn.MaxPool2d(kernel_size=2)
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear')
def forward(self, x):
e1 = self.en_1(x)
e2 = self.en_2(self.maxpool(e1))
e3 = self.en_3(self.maxpool(e2))
e4 = self.en_4(self.maxpool(e3))
d1 = self.de_1(torch.cat([self.upsample(e4), e3], dim=1))
d2 = self.de_2(torch.cat([self.upsample(d1), e2], dim=1))
d3 = self.de_3(torch.cat([self.upsample(d2), e1], dim=1))
d4 = self.de_4(d3)
return d4
# if self.out_channels<2:
# return torch.sigmoid(d4)
# return torch.softmax(d4, 1)
|