Add print of startpos eval

martinnovaak · martinnovaak · commit a59f506f63d3 · 2024-05-10T18:21:18.000+02:00
diff --git a/trainer/model.py b/trainer/model.py
@@ -31,6 +31,39 @@ def forward(self, batch: Batch):
 
         return torch.sigmoid(self.output_layer(hidden_features))
 
+    def eval(self, fen, device):
+        fen = fen.split(" ")[0]
+        stm_features_dense_tensor = torch.zeros(768, device=device)
+        nstm_features_dense_tensor = torch.zeros(768, device=device)
+
+        for rank_idx, rank in enumerate(fen.split('/')):
+            file_idx = 0
+            for char in rank:
+                if char.isdigit():
+                    file_idx += int(char)
+                else:
+                    sq = 8 * (7 - rank_idx) + file_idx
+                    piece_type = {'p': 0, 'n': 1, 'b': 2, 'r': 3, 'q': 4, 'k': 5}[char.lower()]
+
+                    is_black_piece = char.islower()
+                    piece_color = 1 if is_black_piece else 0
+
+                    stm_features_dense_tensor[piece_color * 384 + piece_type * 64 + sq] = 1
+                    nstm_features_dense_tensor[(1 - piece_color) * 384 + piece_type * 64 + (sq ^ 56)] = 1
+
+                    file_idx += 1
+
+        board_stm = stm_features_dense_tensor.to_dense()
+        board_nstm = nstm_features_dense_tensor.to_dense()
+
+        stm_perspective = self.feature_transformer(board_stm)
+        nstm_perspective = self.feature_transformer(board_nstm)
+
+        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()))
+
     def clamp_weights(self):
         self.feature_transformer.weight.data.clamp_(-1.27, 1.27)
         self.output_layer.weight.data.clamp_(-1.27, 1.27)
diff --git a/trainer/train.py b/trainer/train.py
@@ -57,6 +57,8 @@ 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")
 
         optimizer.zero_grad()
