Merge pull request #7 from martinnovaak/logging

martinnovaak · web-flow · commit dbfdd8fda012 · 2024-05-14T15:22:16.000+02:00
Improve logging while training
diff --git a/trainer/model.py b/trainer/model.py
@@ -62,7 +62,7 @@ def eval(self, fen, device):
         hidden_features = torch.cat((stm_perspective, nstm_perspective))
         hidden_features = self.screlu(hidden_features)
 
-        print(int((torch.special.logit(torch.sigmoid(self.output_layer(hidden_features))) * 400).item()))
+        print(self.output_layer(hidden_features) * 400)
 
     def clamp_weights(self):
         self.feature_transformer.weight.data.clamp_(-1.27, 1.27)
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'])
@@ -57,10 +57,12 @@ def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader:
             iterations = 0
             fens = 0
 
-            model.eval("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR", device)
-
             quantize(model, f"network/nnue_{epoch}_scaled.bin")
 
+            save_checkpoint(model, optimizer, epoch, running_loss, "checkpoint.pth")
+
+            model.eval("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR", device)
+
         optimizer.zero_grad()
         prediction = model(batch)
 
@@ -74,4 +76,12 @@ 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)
+
+        if fens % 99_942_400 == 0:
+            print_epoch_stats(epoch, running_loss, iterations, fens, epoch_start_time, time())
+            model.eval("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR", device)
