250 lines
12 KiB
Python
250 lines
12 KiB
Python
import torch
|
|
|
|
|
|
class SimpleCNN(torch.nn.Module):
|
|
def __init__(self,
|
|
in_channels: int = 1,
|
|
out_channels: int = 64,
|
|
hidden_channels: int = 3,
|
|
kernel_size: int = 3,
|
|
dropout_rate: float = 0.01):
|
|
super().__init__()
|
|
layers = []
|
|
for i in range(hidden_channels):
|
|
layers.append(torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2))
|
|
layers.append(torch.nn.ReLU())
|
|
layers.append(torch.nn.BatchNorm2d(num_features=out_channels))
|
|
layers.append(torch.nn.Dropout2d(p=dropout_rate))
|
|
layers.append(torch.nn.MaxPool2d(kernel_size=2, stride=2))
|
|
layers.append(
|
|
torch.nn.ConvTranspose2d(in_channels=out_channels, out_channels=out_channels, kernel_size=2, stride=2))
|
|
in_channels = out_channels
|
|
self.hidden_layers = torch.nn.Sequential(*layers)
|
|
self.out_layer = torch.nn.Conv2d(in_channels, 1, kernel_size=kernel_size, padding=kernel_size // 2)
|
|
self.ac_out = torch.nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
normalized_x = x.float() / 255
|
|
img = self.hidden_layers(normalized_x)
|
|
output = self.out_layer(img)
|
|
scaled_output = self.ac_out(output) * 255
|
|
return scaled_output
|
|
|
|
|
|
# ---------- Past Models ----------
|
|
|
|
# Researched alot about possible CNNs - this is Vanilla UNet Architecture,
|
|
# turns out it doesn't generalize well since its to complex
|
|
class Unet(torch.nn.Module):
|
|
def __init__(self, input_channels: int, padding='same', kernel_size: int = 3):
|
|
super().__init__()
|
|
|
|
# Encoder
|
|
self.max_pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
|
|
self.ac = torch.nn.ReLU()
|
|
# input size = 1
|
|
# input shape torch.Tensor[3, 1, 64, 64]
|
|
print(input_channels)
|
|
self.conv_1 = torch.nn.Conv2d(in_channels=1,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_1_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_1 = torch.nn.BatchNorm2d(input_channels)
|
|
# apply conv layer 2 times, ReLu after every application, max pooling after every 2 applications of conv layer
|
|
input_channels *= 2
|
|
# input size = 128
|
|
self.conv_2 = torch.nn.Conv2d(in_channels=input_channels // 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_2_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_2 = torch.nn.BatchNorm2d(input_channels)
|
|
input_channels *= 2
|
|
# input size = 256
|
|
self.conv_3 = torch.nn.Conv2d(in_channels=input_channels // 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_3_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_3 = torch.nn.BatchNorm2d(input_channels)
|
|
input_channels *= 2
|
|
# Bridge
|
|
# input size = 512
|
|
self.conv_4 = torch.nn.Conv2d(in_channels=input_channels // 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_4_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_4 = torch.nn.BatchNorm2d(input_channels)
|
|
# Decoder
|
|
self.up_samp = torch.nn.UpsamplingBilinear2d(scale_factor=2)
|
|
# input size = 256, kernel size = 3 - 1
|
|
input_channels //= 2
|
|
self.conv_5 = torch.nn.Conv2d(in_channels=input_channels * 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
# concatenate in forward function (tensor after last application of conv 3 merged with current one(up sampled))
|
|
# input size = 512
|
|
self.conv_6 = torch.nn.Conv2d(in_channels=input_channels * 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_6_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_5 = torch.nn.BatchNorm2d(input_channels)
|
|
# apply conv_6 2 times
|
|
# up sampling
|
|
# input size = 128, kernel size = 3 - 1
|
|
input_channels //= 2
|
|
self.conv_7 = torch.nn.Conv2d(in_channels=input_channels * 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
# concatenate in forward function (tensor after conv 2 with current one)
|
|
# input size = 128, kernel size = 3
|
|
self.conv_8 = torch.nn.Conv2d(in_channels=input_channels * 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_8_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_6 = torch.nn.BatchNorm2d(input_channels)
|
|
# up sampling
|
|
# input size = 64, kernel size = 3 - 1
|
|
input_channels //= 2
|
|
self.conv_9 = torch.nn.Conv2d(in_channels=input_channels * 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
# concatenate in forward function (tensor after conv1 with current one)
|
|
self.conv_10 = torch.nn.Conv2d(in_channels=input_channels * 2,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv_10_1 = torch.nn.Conv2d(in_channels=input_channels,
|
|
out_channels=input_channels,
|
|
kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.batch_norm_7 = torch.nn.BatchNorm2d(input_channels)
|
|
for m in self.modules():
|
|
if isinstance(m, torch.nn.Conv2d):
|
|
torch.nn.init.xavier_uniform_(m.weight)
|
|
self.out = torch.nn.Conv2d(64, 1, kernel_size=1)
|
|
self.ac_out = torch.nn.Sigmoid()
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
# Encoder
|
|
x = x.float() / 255
|
|
img = self.ac(self.conv_1(x))
|
|
img = self.batch_norm_1(img)
|
|
conv1 = img = self.ac(self.conv_1_1(img))
|
|
img = self.batch_norm_1(img)
|
|
img = self.max_pool(img)
|
|
img = self.ac(self.conv_2(img))
|
|
img = self.batch_norm_2(img)
|
|
conv2 = img = self.ac(self.conv_2_1(img))
|
|
img = self.batch_norm_2(img)
|
|
img = self.max_pool(img)
|
|
img = self.ac(self.conv_3(img))
|
|
img = self.batch_norm_3(img)
|
|
conv3 = img = self.ac(self.conv_3_1(img))
|
|
img = self.batch_norm_3(img)
|
|
img = self.max_pool(img)
|
|
# Bridge
|
|
img = self.ac(self.conv_4(img))
|
|
img = self.batch_norm_4(img)
|
|
img = torch.nn.Dropout(p=0.001)(img)
|
|
img = self.ac(self.conv_4_1(img))
|
|
img = self.batch_norm_4(img)
|
|
img = torch.nn.Dropout(p=0.001)(img)
|
|
# Decoder
|
|
img = self.up_samp(img)
|
|
img = self.ac(self.conv_5(img))
|
|
img = self.batch_norm_5(img)
|
|
img = torch.cat([conv3, img], dim=1)
|
|
img = self.ac(self.conv_6(img))
|
|
img = self.batch_norm_5(img)
|
|
img = self.ac(self.conv_6_1(img))
|
|
img = self.batch_norm_5(img)
|
|
|
|
img = self.up_samp(img)
|
|
img = self.ac(self.conv_7(img))
|
|
img = self.batch_norm_6(img)
|
|
img = torch.concatenate([conv2, img], dim=1)
|
|
img = self.ac(self.conv_8(img))
|
|
img = self.batch_norm_6(img)
|
|
img = self.ac(self.conv_8_1(img))
|
|
img = self.batch_norm_6(img)
|
|
|
|
img = self.up_samp(img)
|
|
img = self.ac(self.conv_9(img))
|
|
img = self.batch_norm_7(img)
|
|
img = torch.cat([conv1, img], dim=1)
|
|
img = self.ac(self.conv_10(img))
|
|
img = self.batch_norm_7(img)
|
|
img = self.ac(self.conv_10_1(img))
|
|
img = self.batch_norm_7(img)
|
|
|
|
# Output
|
|
output = self.ac_out(self.out(img))
|
|
return output
|
|
|
|
|
|
# Testing around with basic idea of UNet above but decreasing model complexity - did not work to well either
|
|
class Model(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.in_chn = torch.nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, padding=1)
|
|
self.bn2d_64 = torch.nn.BatchNorm2d(num_features=64)
|
|
self.ac = torch.nn.ReLU()
|
|
self.max_pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
|
|
self.conv_64_128 = torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
|
|
self.bn2d_128 = torch.nn.BatchNorm2d(num_features=128)
|
|
self.conv_128_256 = torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1)
|
|
self.bn2d_256 = torch.nn.BatchNorm2d(num_features=256)
|
|
|
|
self.conv_256_512 = torch.nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1)
|
|
self.bn2d_512 = torch.nn.BatchNorm2d(num_features=512)
|
|
|
|
# concat
|
|
self.conv_256_128 = torch.nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=1)
|
|
self.conv_128_64 = torch.nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1)
|
|
self.out_chn = torch.nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
|
|
self.ac_out = torch.nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
img = self.bn2d_64(self.ac(self.in_chn(x)))
|
|
state_1 = img = self.bn2d_128(self.ac(self.conv_64_128(img)))
|
|
img = torch.nn.Dropout(p=0.08)(img)
|
|
img = self.bn2d_256(self.ac(self.conv_128_256(img)))
|
|
img = torch.nn.Dropout(p=0.01)(img)
|
|
img = self.bn2d_256(self.ac(self.conv_128_256(img)))
|
|
img = torch.nn.Dropout(p=0.01)(img)
|
|
img = self.bn2d_128(self.ac(self.conv_256_128(img)))
|
|
img = torch.nn.Dropout(p=0.1)(img)
|
|
img = torch.concat([state_1, img], dim=1)
|
|
img = self.bn2d_128(self.ac(self.conv_256_128(img)))
|
|
img = torch.nn.Dropout(p=0.01)(img)
|
|
img = self.bn2d_64(self.ac(self.conv_128_64(img)))
|
|
img = torch.nn.Dropout(p=0.08)(img)
|
|
img = self.out_chn(img)
|
|
return self.ac_out(img)
|