some cleanup to trainer

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -146,18 +146,26 @@ def __init__(
 
         self.clip_grad_norm = clip_grad_norm
 
+        # steps
+
+        self.steps = 0
+
     @property
     def is_main(self):
         return self.fabric.global_rank == 0
 
+    def print(self, *args, **kwargs):
+        self.fabric.print(*args, **kwargs)
+
+    def log(self, **log_data):
+        self.fabric.log_dict(log_data, step = self.steps)
+
     def __call__(
         self
     ):
-        dl = iter(self.dataloader)
-
-        steps = 0
+        dl = cycle(self.dataloader)
 
-        while steps < self.num_train_steps:
+        while self.steps < self.num_train_steps:
 
             for grad_accum_step in range(self.grad_accum_every):
                 is_accumulating = grad_accum_step < (self.grad_accum_every - 1)
@@ -169,7 +177,9 @@ def __call__(
 
                     self.fabric.backward(loss / self.grad_accum_every)
 
-            print(f'loss: {loss.item():.3f}')
+            self.log(loss = loss)
+
+            self.print(f'loss: {loss.item():.3f}')
 
             self.fabric.clip_gradients(self.model, self.optimizer, max_norm = self.clip_grad_norm)
 
@@ -181,6 +191,6 @@ def __call__(
             self.scheduler.step()
             self.optimizer.zero_grad()
 
-            steps += 1
+            self.steps += 1
 
         print(f'training complete')

pyproject.toml

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