Snapshot

interposer.py

@@ -1,56 +1,57 @@
 import torch
 import torch.nn as nn
-import numpy as np
 
-class Block(nn.Module):
-	def __init__(self, size):
+class ResBlock(nn.Module):
+	"""Block with residuals"""
+	def __init__(self, ch):
 		super().__init__()
 		self.join = nn.ReLU()
 		self.long = nn.Sequential(
-			nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),
+			nn.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1),
 			nn.LeakyReLU(0.1),
-			nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),
+			nn.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1),
 			nn.LeakyReLU(0.1),
-			nn.Conv2d(size, size, kernel_size=3, stride=1, padding=1),
-			nn.Dropout(0.2)
+			nn.Conv2d(ch, ch, kernel_size=3, stride=1, padding=1),
 		)
 	def forward(self, x):
-		y = self.long(x)
-		z = self.join(y + x)
-		return z
+		return self.join(self.long(x) + x)
 
-class Interposer(nn.Module):
-	def __init__(self):
+class ExtractBlock(nn.Module):
+	"""Increase no. of channels by [out/in]"""
+	def __init__(self, ch_in, ch_out):
 		super().__init__()
-		self.chan = 4       # in/out channels
-		self.hid = 128
+		self.join  = nn.ReLU()
+		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.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1),
+			nn.LeakyReLU(0.1),
+			nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1),
+			nn.Dropout(0.1)
+		)
+	def forward(self, x):
+		return self.join(self.long(x) + self.short(x))
 
-		# expand channels
-		self.head_join  = nn.ReLU()
-		self.head_short = nn.Conv2d(self.chan, self.hid, kernel_size=3, stride=1, padding=1)
-		self.head_long  = nn.Sequential(
-			nn.Conv2d(self.chan, self.hid, kernel_size=3, stride=1, padding=1),
-			nn.LeakyReLU(0.1),
-			nn.Conv2d(self.hid,  self.hid, kernel_size=3, stride=1, padding=1),
-			nn.LeakyReLU(0.1),
-			nn.Conv2d(self.hid,  self.hid, kernel_size=3, stride=1, padding=1),
-		)
-		# not sure if this is how residuals work
+class InterposerModel(nn.Module):
+	"""Main neural network"""
+	def __init__(self, ch_in=4, ch_out=4, ch_mid=64, scale=1.0):
+		super().__init__()
+		self.scale  = scale
+		self.ch_in  = ch_in
+		self.ch_out = ch_out
+		self.ch_mid = ch_mid
+
+		self.head = ExtractBlock(self.ch_in, self.ch_mid)
 		self.core = nn.Sequential(
-			Block(self.hid),
-			Block(self.hid),
-			Block(self.hid),
-		)
-		# reduce channels
-		self.tail = nn.Sequential(
-			nn.ReLU(),
-			nn.Conv2d(self.hid, self.chan, kernel_size=3, stride=1, padding=1)
+			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),
 		)
+		self.tail = nn.Conv2d(self.ch_mid, self.ch_out, kernel_size=3, stride=1, padding=1)
 
 	def forward(self, x):
-		y = self.head_join(
-			self.head_long(x)+
-			self.head_short(x)
-		)
+		y = self.head(x)
 		z = self.core(y)
 		return self.tail(z)