Rename ArcFace Converter to Embedding Converter, Add EmbeddingDataset, Add learning rate to config

2026-07-28 16:08:50 +02:00 · 2025-03-11 14:43:06 +01:00
parent 62a69cddd2
commit 1b6e7a6ca5
23 changed files with 70 additions and 66 deletions
@@ -0,0 +1,22 @@
+import configparser
+from os import makedirs
+
+import torch
+
+from .training import EmbeddingConverterTrainer
+
+CONFIG = configparser.ConfigParser()
+CONFIG.read('config.ini')
+
+
+def export() -> None:
+	directory_path = CONFIG.get('exporting', 'directory_path')
+	source_path = CONFIG.get('exporting', 'source_path')
+	target_path = CONFIG.get('exporting', 'target_path')
+	opset_version = CONFIG.getint('exporting', 'opset_version')
+
+	makedirs(directory_path, exist_ok = True)
+	embedding_converter_trainer = EmbeddingConverterTrainer.load_from_checkpoint(source_path, map_location = 'cpu')
+	embedding_converter_trainer.eval()
+	input_tensor = torch.randn(1, 512)
+	torch.onnx.export(embedding_converter_trainer, input_tensor, target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = opset_version)
@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+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.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
@@ -0,0 +1,79 @@
+import configparser
+from os import makedirs
+from os.path import isfile
+from typing import List
+
+import numpy
+numpy.bool = numpy.bool_
+from mxnet.io import ImageRecordIter
+from onnxruntime import InferenceSession
+from tqdm import tqdm
+
+from .types import Embedding, EmbeddingDataset, VisionFrame
+
+CONFIG = configparser.ConfigParser()
+CONFIG.read('config.ini')
+
+
+def prepare_crop_vision_frame(crop_vision_frame : VisionFrame) -> VisionFrame:
+	crop_vision_frame = crop_vision_frame.astype(numpy.float32) / 255.0
+	crop_vision_frame = (crop_vision_frame - 0.5) * 2
+	return crop_vision_frame
+
+
+def create_inference_session(model_path : str, execution_providers : List[str]) -> InferenceSession:
+	inference_session = InferenceSession(model_path, providers = execution_providers)
+	return inference_session
+
+
+def forward(inference_session : InferenceSession, crop_vision_frame : VisionFrame) -> Embedding:
+	embedding = inference_session.run(None,
+	{
+		'input': crop_vision_frame
+	})[0]
+
+	return embedding
+
+
+def create_embedding_dataset(dataset_reader : ImageRecordIter, source_inference_session : InferenceSession, target_inference_session : InferenceSession) -> EmbeddingDataset:
+	dataset_process_limit = CONFIG.getint('preparing.dataset', 'process_limit')
+	embedding_pairs = []
+
+	with tqdm(total = dataset_process_limit) as progress:
+		for batch in dataset_reader:
+			crop_vision_frame = batch.data[0].asnumpy()
+			crop_vision_frame = prepare_crop_vision_frame(crop_vision_frame)
+			source_embedding = forward(source_inference_session, crop_vision_frame)
+			target_embedding = forward(target_inference_session, crop_vision_frame)
+			embedding_pairs.append([ source_embedding, target_embedding ])
+			progress.update()
+
+			if progress.n == dataset_process_limit:
+				return numpy.concatenate(embedding_pairs, axis = 1).T
+
+	return numpy.concatenate(embedding_pairs, axis = 1).T
+
+
+def prepare() -> None:
+	dataset_path = CONFIG.get('preparing.dataset', 'dataset_path')
+	dataset_crop_size = CONFIG.getint('preparing.dataset', 'crop_size')
+	model_source_path = CONFIG.get('preparing.model', 'source_path')
+	model_target_path = CONFIG.get('preparing.model', 'target_path')
+	input_directory_path = CONFIG.get('preparing.input', 'directory_path')
+	input_source_path = CONFIG.get('preparing.input', 'source_path')
+	input_target_path = CONFIG.get('preparing.input', 'target_path')
+	execution_providers = CONFIG.get('execution', 'providers').split(' ')
+
+	makedirs(input_directory_path, exist_ok = True)
+	if isfile(dataset_path) and isfile(model_source_path) and isfile(model_target_path):
+		dataset_reader = ImageRecordIter(
+			path_imgrec = dataset_path,
+			data_shape = (3, dataset_crop_size, dataset_crop_size),
+			batch_size = 1,
+			shuffle = False
+		)
+		source_inference_session = create_inference_session(model_source_path, execution_providers)
+		target_inference_session = create_inference_session(model_target_path, execution_providers)
+		embedding_dataset = create_embedding_dataset(dataset_reader, source_inference_session, target_inference_session)
+		numpy.save(input_source_path, embedding_dataset[..., 0].T)
+		numpy.save(input_target_path, embedding_dataset[..., 1].T)
@@ -0,0 +1,116 @@
+import configparser
+from typing import Any, Tuple
+
+import numpy
+import pytorch_lightning
+import torch
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.tuner.tuning import Tuner
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset, TensorDataset, random_split
+
+from .models.embedding_converter import EmbeddingConverter
+from .types import Batch, Loader
+
+CONFIG = configparser.ConfigParser()
+CONFIG.read('config.ini')
+
+
+class EmbeddingConverterTrainer(pytorch_lightning.LightningModule):
+	def __init__(self) -> None:
+		super(EmbeddingConverterTrainer, self).__init__()
+		self.embedding_converter = EmbeddingConverter()
+		self.mse_loss = torch.nn.MSELoss()
+
+	def forward(self, source_embedding : Tensor) -> Tensor:
+		return self.embedding_converter(source_embedding)
+
+	def training_step(self, batch : Batch, batch_index : int) -> Tensor:
+		source, target = batch
+		output = self(source)
+		loss_training = self.mse_loss(output, target)
+		self.log('loss_training', loss_training, prog_bar = True)
+		return loss_training
+
+	def validation_step(self, batch : Batch, batch_index : int) -> Tensor:
+		source, target = batch
+		output = self(source)
+		loss_validation = self.mse_loss(output, target)
+		self.log('loss_validation', loss_validation, prog_bar = True)
+		return loss_validation
+
+	def configure_optimizers(self) -> Any:
+		learning_rate = CONFIG.getfloat('training.trainer', 'learning_rate')
+		optimizer = torch.optim.Adam(self.parameters(), lr = learning_rate)
+		scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
+
+		return\
+		{
+			'optimizer': optimizer,
+			'lr_scheduler':
+			{
+				'scheduler': scheduler,
+				'monitor': 'train_loss',
+				'interval': 'epoch',
+				'frequency': 1
+			}
+		}
+
+
+def create_loaders() -> Tuple[Loader, Loader]:
+	loader_batch_size = CONFIG.getint('training.loader', 'batch_size')
+	loader_num_workers = CONFIG.getint('training.loader', 'num_workers')
+
+	training_dataset, validate_dataset = split_dataset()
+	training_loader = DataLoader(training_dataset, batch_size = loader_batch_size, num_workers = loader_num_workers, shuffle = True, pin_memory = True)
+	validation_loader = DataLoader(validate_dataset, batch_size = loader_batch_size, num_workers = loader_num_workers, shuffle = False, pin_memory = True)
+	return training_loader, validation_loader
+
+
+def split_dataset() -> Tuple[Dataset[Any], Dataset[Any]]:
+	input_source_path = CONFIG.get('preparing.input', 'source_path')
+	input_target_path = CONFIG.get('preparing.input', 'target_path')
+	loader_split_ratio = CONFIG.getfloat('training.loader', 'split_ratio')
+
+	source_input = torch.from_numpy(numpy.load(input_source_path)).float()
+	target_input = torch.from_numpy(numpy.load(input_target_path)).float()
+	dataset = TensorDataset(source_input, target_input)
+
+	dataset_size = len(dataset)
+	training_size = int(loader_split_ratio * len(dataset))
+	validation_size = int(dataset_size - training_size)
+	training_dataset, validate_dataset = random_split(dataset, [ training_size, validation_size ])
+	return training_dataset, validate_dataset
+
+
+def create_trainer() -> Trainer:
+	trainer_max_epochs = CONFIG.getint('training.trainer', 'max_epochs')
+	output_directory_path = CONFIG.get('training.output', 'directory_path')
+	output_file_pattern = CONFIG.get('training.output', 'file_pattern')
+
+	return Trainer(
+		max_epochs = trainer_max_epochs,
+		callbacks =
+		[
+			ModelCheckpoint(
+				monitor = 'train_loss',
+				dirpath = output_directory_path,
+				filename = output_file_pattern,
+				every_n_epochs = 10,
+				save_top_k = 3,
+				save_last = True
+			)
+		],
+		enable_progress_bar = True,
+		log_every_n_steps = 2
+	)
+
+
+def train() -> None:
+	trainer = create_trainer()
+	training_loader, validation_loader = create_loaders()
+	embedding_converter = EmbeddingConverterTrainer()
+	tuner = Tuner(trainer)
+	tuner.lr_find(embedding_converter, training_loader, validation_loader)
+	trainer.fit(embedding_converter, training_loader, validation_loader)
@@ -0,0 +1,14 @@
+from typing import Any, Tuple
+
+from numpy.typing import NDArray
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+Batch = Tuple[Tensor, Tensor]
+Loader = DataLoader[Tuple[Tensor, ...]]
+
+Embedding = NDArray[Any]
+EmbeddingDataset = NDArray[Embedding]
+FaceLandmark5 = NDArray[Any]
+
+VisionFrame = NDArray[Any]