CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-07-28 16:08:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Improve naming for arcface converter trainer

Browse Source
This commit is contained in:
henryruhs
2025-01-15 09:26:21 +01:00
parent abce63007d
commit f3e1c3feaa
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
arcface_converter/src/training.py
+6 -6
View File
@@ -30,16 +30,16 @@ class ArcFaceConverterTrainer(pytorch_lightning.LightningModule):
return self.model(source_embedding)
def training_step(self, batch : Batch, batch_index : int) -> Tensor:
source, target = batch
output = self(source)
loss = self.loss_fn(output, target)
source_embedding, target_embedding = batch
output_embedding = self(source_embedding)
loss = self.loss_fn(output_embedding, target_embedding)
self.log('train_loss', loss, prog_bar = True, logger = True)
return loss
def validation_step(self, batch : Batch, batch_index : int) -> Tensor:
source, target = batch
output = self(source)
loss = self.loss_fn(output, target)
source_embedding, target_embedding = batch
output_embedding = self(source_embedding)
loss = self.loss_fn(output_embedding, target_embedding)
self.log('val_loss', loss, prog_bar = True, logger = True)
return loss