Use sequential whenever possible

face_swapper/src/networks/aad.py

@@ -162,10 +162,16 @@ class FeatureModulation(nn.Module):
 class PixelShuffleUpSample(nn.Module):
 	def __init__(self, input_channels : int, output_channels : int) -> None:
 		super().__init__()
-		self.conv = nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1)
-		self.pixel_shuffle = nn.PixelShuffle(upscale_factor = 2)
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1),
+			nn.PixelShuffle(upscale_factor = 2)
+		)
 
 	def forward(self, input_tensor : Tensor) -> Tensor:
-		output_tensor = self.conv(input_tensor.view(input_tensor.shape[0], -1, 1, 1))
-		output_tensor = self.pixel_shuffle(output_tensor)
+		temp_tensor = input_tensor.view(input_tensor.shape[0], -1, 1, 1)
+		output_tensor = self.sequences(temp_tensor)
 		return output_tensor

face_swapper/src/networks/masknet.py

@@ -21,7 +21,7 @@ class MaskNet(nn.Module):
 		return nn.ModuleList(
 		[
 			DownSample(input_channels, num_filters),
-			DownSample(num_filters, num_filters * 2),
+			DownSample(num_filters, num_filters * 2)
 		])
 
 	@staticmethod
@@ -74,26 +74,34 @@ class BottleNeck(nn.Module):
 class UpSample(nn.Module):
 	def __init__(self, input_channels : int, output_channels : int) -> None:
 		super().__init__()
-		self.conv_transpose = nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 2, stride = 2)
-		self.relu = nn.ReLU(inplace = True)
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 2, stride = 2),
+			nn.ReLU(inplace = True)
+		)
 
 	def forward(self, input_tensor : Tensor) -> Tensor:
-		output_tensor = self.conv_transpose(input_tensor)
-		output_tensor = self.relu(output_tensor)
+		output_tensor = self.sequences(input_tensor)
 		return output_tensor
 
 
 class DownSample(nn.Module):
 	def __init__(self, input_channels : int, output_channels : int) -> None:
 		super().__init__()
-		self.conv = nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1, bias = False)
-		self.batch_norm = nn.BatchNorm2d(output_channels)
-		self.relu = nn.ReLU(inplace = True)
-		self.max_pool = nn.MaxPool2d(2)
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(input_channels, output_channels, kernel_size = 3, padding = 1, bias = False),
+			nn.BatchNorm2d(output_channels),
+			nn.ReLU(inplace = True),
+			nn.MaxPool2d(2)
+		)
 
 	def forward(self, input_tensor : Tensor) -> Tensor:
-		output_tensor = self.conv(input_tensor)
-		output_tensor = self.batch_norm(output_tensor)
-		output_tensor = self.relu(output_tensor)
-		output_tensor = self.max_pool(output_tensor)
+		output_tensor = self.sequences(input_tensor)
 		return output_tensor

face_swapper/src/networks/unet.py

@@ -103,14 +103,18 @@ class UNet(nn.Module):
 class UpSample(nn.Module):
 	def __init__(self, input_channels : int, output_channels : int) -> None:
 		super().__init__()
-		self.conv_transpose = nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
-		self.batch_norm = nn.BatchNorm2d(output_channels)
-		self.leaky_relu = nn.LeakyReLU(0.1, inplace = True)
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.ConvTranspose2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False),
+			nn.BatchNorm2d(output_channels),
+			nn.LeakyReLU(0.1, inplace = True)
+		)
 
 	def forward(self, input_tensor : Tensor, skip_tensor : Tensor) -> Tensor:
-		output_tensor = self.conv_transpose(input_tensor)
-		output_tensor = self.batch_norm(output_tensor)
-		output_tensor = self.leaky_relu(output_tensor)
+		output_tensor = self.sequences(input_tensor)
 		output_tensor = torch.cat((output_tensor, skip_tensor), dim = 1)
 		return output_tensor
 
@@ -118,12 +122,16 @@ class UpSample(nn.Module):
 class DownSample(nn.Module):
 	def __init__(self, input_channels : int, output_channels : int) -> None:
 		super().__init__()
-		self.conv = nn.Conv2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
-		self.batch_norm = nn.BatchNorm2d(output_channels)
-		self.leaky_relu = nn.LeakyReLU(0.1, inplace = True)
+		self.sequences = self.create_sequences(input_channels, output_channels)
+
+	@staticmethod
+	def create_sequences(input_channels : int, output_channels : int) -> nn.Sequential:
+		return nn.Sequential(
+			nn.Conv2d(input_channels, output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False),
+			nn.BatchNorm2d(output_channels),
+			nn.LeakyReLU(0.1, inplace = True)
+		)
 
 	def forward(self, input_tensor : Tensor) -> Tensor:
-		output_tensor = self.conv(input_tensor)
-		output_tensor = self.batch_norm(output_tensor)
-		output_tensor = self.leaky_relu(output_tensor)
+		output_tensor = self.sequences(input_tensor)
 		return output_tensor