Modernize to use ModuleList, Fix some types

embedding_converter/src/models/embedding_converter.py

@@ -1,21 +1,31 @@
 import torch
 import torch.nn as nn
-from torch import Tensor
+
+from embedding_converter.src.types import VisionTensor
 
 
 class EmbeddingConverter(nn.Module):
 	def __init__(self) -> None:
 		super(EmbeddingConverter, self).__init__()
-		self.fc1 = nn.Linear(512, 1024)
-		self.fc2 = nn.Linear(1024, 2048)
-		self.fc3 = nn.Linear(2048, 1024)
-		self.fc4 = nn.Linear(1024, 512)
+		self.layers = self.create_layers()
 		self.activation = nn.LeakyReLU()
 
-	def forward(self, inputs : Tensor) -> Tensor:
-		norm_inputs = inputs / torch.norm(inputs)
-		outputs = self.activation(self.fc1(norm_inputs))
-		outputs = self.activation(self.fc2(outputs))
-		outputs = self.activation(self.fc3(outputs))
-		outputs = self.fc4(outputs)
-		return outputs
+	@staticmethod
+	def create_layers() -> nn.ModuleList:
+		layers = nn.ModuleList(
+		[
+			nn.Linear(512, 1024),
+			nn.Linear(1024, 2048),
+			nn.Linear(2048, 1024),
+			nn.Linear(1024, 512)
+		])
+		return layers
+
+	def forward(self, input_tensor: VisionTensor) -> VisionTensor:
+		output_tensor = input_tensor / torch.norm(input_tensor)
+
+		for layer in self.layers[:-1]:
+			output_tensor = self.activation(layer(output_tensor))
+
+		output_tensor = self.layers[-1](output_tensor)
+		return output_tensor