log the entire loss breakdown

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1757,9 +1757,9 @@ def forward(
 # https://arxiv.org/abs/2206.00364
 
 class DiffusionLossBreakdown(NamedTuple):
-    mse: Float['']
-    bond: Float['']
-    smooth_lddt: Float['']
+    diffusion_mse: Float['']
+    diffusion_bond: Float['']
+    diffusion_smooth_lddt: Float['']
 
 class ElucidatedAtomDiffusionReturn(NamedTuple):
     loss: Float['']
@@ -2618,14 +2618,17 @@ def forward(
 # main class
 
 class LossBreakdown(NamedTuple):
-    diffusion: Float['']
+    total_loss: Float['']
+    total_diffusion: Float['']
     distogram: Float['']
     pae: Float['']
     pde: Float['']
     plddt: Float['']
     resolved: Float['']
     confidence: Float['']
-    diffusion_loss_breakdown: DiffusionLossBreakdown
+    diffusion_mse: Float['']
+    diffusion_bond: Float['']
+    diffusion_smooth_lddt: Float['']
 
 class Alphafold3(Module):
     """ Algorithm 1 """
@@ -3215,14 +3218,15 @@ def forward(
             return loss
 
         loss_breakdown = LossBreakdown(
+            total_loss = loss,
+            total_diffusion = diffusion_loss,
             pae = pae_loss,
             pde = pde_loss,
             plddt = plddt_loss,
             resolved = resolved_loss,
             distogram = distogram_loss,
-            diffusion = diffusion_loss,
             confidence = confidence_loss,
-            diffusion_loss_breakdown = diffusion_loss_breakdown
+            **diffusion_loss_breakdown._asdict()
         )
 
         return loss, loss_breakdown

alphafold3_pytorch/trainer.py

@@ -173,11 +173,14 @@ def __call__(
                 inputs = next(dl)
 
                 with self.fabric.no_backward_sync(self.model, enabled = is_accumulating):
-                    loss = self.model(**inputs)
+                    loss, loss_breakdown = self.model(
+                        **inputs,
+                        return_loss_breakdown = True
+                    )
 
                     self.fabric.backward(loss / self.grad_accum_every)
 
-            self.log(loss = loss)
+            self.log(**loss_breakdown._asdict())
 
             self.print(f'loss: {loss.item():.3f}')
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.0.45"
+version = "0.0.46"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }