Version 4

interposer.py

@@ -6,14 +6,17 @@ class ResBlock(nn.Module):
 	def __init__(self, ch):
 		super().__init__()
 		self.join = nn.ReLU()
+		self.norm = nn.BatchNorm2d(ch)
 		self.long = nn.Sequential(
 			nn.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1),
-			nn.LeakyReLU(0.1),
+			nn.SiLU(),
 			nn.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1),
-			nn.LeakyReLU(0.1),
+			nn.SiLU(),
 			nn.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1),
+			nn.Dropout(0.1)
 		)
 	def forward(self, x):
+		x = self.norm(x)
 		return self.join(self.long(x) + x)
 
 class ExtractBlock(nn.Module):
@@ -24,9 +27,9 @@ class ExtractBlock(nn.Module):
 		self.short = nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1)
 		self.long  = nn.Sequential(
 			nn.Conv2d( ch_in, ch_out, kernel_size=3, stride=1, padding=1),
-			nn.LeakyReLU(0.1),
+			nn.SiLU(),
 			nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1),
-			nn.LeakyReLU(0.1),
+			nn.SiLU(),
 			nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1),
 			nn.Dropout(0.1)
 		)
@@ -35,19 +38,20 @@ class ExtractBlock(nn.Module):
 
 class InterposerModel(nn.Module):
 	"""Main neural network"""
-	def __init__(self, ch_in=4, ch_out=4, ch_mid=64, scale=1.0):
+	def __init__(self, ch_in=4, ch_out=4, ch_mid=64, scale=1.0, blocks=12):
 		super().__init__()
-		self.scale  = scale
 		self.ch_in  = ch_in
 		self.ch_out = ch_out
 		self.ch_mid = ch_mid
+		self.blocks = blocks
+		self.scale  = scale
 
 		self.head = ExtractBlock(self.ch_in, self.ch_mid)
 		self.core = nn.Sequential(
 			nn.Upsample(scale_factor=self.scale, mode="nearest"),
-			ResBlock(self.ch_mid), ResBlock(self.ch_mid), ResBlock(self.ch_mid), ResBlock(self.ch_mid),
-			ResBlock(self.ch_mid), ResBlock(self.ch_mid), ResBlock(self.ch_mid), ResBlock(self.ch_mid),
-			ResBlock(self.ch_mid), ResBlock(self.ch_mid), ResBlock(self.ch_mid), ResBlock(self.ch_mid),
+			*[ResBlock(self.ch_mid) for _ in range(blocks)],
+			nn.BatchNorm2d(self.ch_mid),
+			nn.SiLU(),
 		)
 		self.tail = nn.Conv2d(self.ch_mid, self.ch_out, kernel_size=3, stride=1, padding=1)