Merge branch 'main' into refactor

martinnovaak · web-flow · commit 98329fabc595 · 2024-08-08T22:42:26.000+02:00
diff --git a/trainer/main.py b/trainer/main.py
@@ -7,6 +7,7 @@
 from batchloader import BatchLoader
 from model import PerspectiveNetwork
 from train import train
+from quantize import load_quantized_net
 
 
 def main():
@@ -31,6 +32,8 @@ def load_config(config_path="config.json"):
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     model = PerspectiveNetwork(config["hidden_layer_size"]).to(device)
+    #model = load_quantized_net("nnue.bin", config["hidden_layer_size"], 403, 64).to(device)
+    model.eval("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR", device)
 
     paths = [os.path.join(data_root.encode("utf-8"), file.encode("utf-8")) for file in os.listdir(data_root)]
 
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/quantize.py b/trainer/quantize.py
@@ -2,6 +2,8 @@
 import struct
 import numpy as np
 
+from model import PerspectiveNetwork
+
 QA = 403
 QB = 64
 QAB = QA * QB
@@ -39,4 +41,4 @@ def load_quantized_net(bin_path, hl_size, qa, qb):
     model.output_layer.weight.data = torch.tensor(np.array(output_weights).reshape(1, 2 * hl_size) / qb, dtype=torch.float32)
     model.output_layer.bias.data = torch.tensor(np.array(output_bias) / (qa * qb), dtype=torch.float32)
 
-    return model
+    return model
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)
