Improve printing

martinnovaak · martinnovaak · commit ebde75181f17 · 2024-05-11T00:18:20.000+02:00
diff --git a/trainer/train.py b/trainer/train.py
@@ -7,7 +7,7 @@
 
 def print_epoch_stats(epoch, running_loss, iterations, fens, start_time, current_time):
     epoch_time = current_time - start_time
-    message = ("\nepoch {:<2} | time: {:.2f} s | epoch loss: {:.4f} | speed: {:.2f} pos/s"
+    message = ("\nepoch {:<2} | time: {:.2f} s | epoch loss: {:.7f} | speed: {:.2f} pos/s"
                .format(epoch, epoch_time, running_loss.item() / iterations, fens / epoch_time))
     print(message)
 
@@ -20,7 +20,7 @@ def save_checkpoint(model, optimizer, epoch, loss, filename):
     }
     torch.save(checkpoint, filename)
 
-def load_checkpoint(model, optimizer, filename, resume_training=False):
+def load_checkpoint(model, optimizer, filename):
     checkpoint = torch.load(filename)
     model.load_state_dict(checkpoint['model_state_dict'])
     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
@@ -61,6 +61,8 @@ def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader:
 
             quantize(model, f"network/nnue_{epoch}_scaled.bin")
 
+            save_checkpoint(model, optimizer, epoch, running_loss, "checkpoint.pth")
+
         optimizer.zero_grad()
         prediction = model(batch)
 
@@ -74,4 +76,10 @@ def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader:
         fens += batch.size
 
         if fens % 163_840 == 0:
-            print("\rTotal fens parsed in this epoch:", fens, end='', flush=True)
+            epoch_time = time() - epoch_start_time
+            formatted_fens = "{0:_}".format(fens)
+            formatted_speed = "{0:_}".format(int(fens / epoch_time))
+            print("\rTotal fens parsed in this epoch: {}, Time: {:.2f} s, Speed: {} pos/s".format(formatted_fens,
+                                                                                                  epoch_time,
+                                                                                                  formatted_speed),
+                  end='', flush=True)
