Add FeatureSet and Model Type Hints (official-stockfish#343)

Disservin · web-flow · commit 4b1477a992f3 · 2025-09-21T16:06:11.000+02:00
diff --git a/cross_check_eval.py b/cross_check_eval.py
@@ -8,9 +8,10 @@
 import serialize
 import data_loader
 from model import NNUE
+from features.feature_set import FeatureSet
 
 
-def read_model(nnue_path, feature_set):
+def read_model(nnue_path, feature_set: FeatureSet):
     with open(nnue_path, "rb") as f:
         reader = serialize.NNUEReader(f, feature_set)
         return reader.model
diff --git a/data_loader/dataset.py b/data_loader/dataset.py
@@ -57,7 +57,7 @@ def __del__(self):
 class TrainingDataProvider:
     def __init__(
         self,
-        feature_set,
+        feature_set: str,
         create_stream,
         destroy_stream,
         fetch_next,
@@ -113,7 +113,7 @@ def __del__(self):
 class SparseBatchProvider(TrainingDataProvider):
     def __init__(
         self,
-        feature_set,
+        feature_set: str,
         filenames,
         batch_size,
         cyclic=True,
@@ -137,7 +137,7 @@ def __init__(
 class SparseBatchDataset(torch.utils.data.IterableDataset):
     def __init__(
         self,
-        feature_set,
+        feature_set: str,
         filenames,
         batch_size,
         cyclic=True,
diff --git a/data_loader/stream.py b/data_loader/stream.py
@@ -2,6 +2,7 @@
 
 from ._native import c_lib
 from .config import CDataloaderSkipConfig, DataloaderSkipConfig
+from features.feature_set import FeatureSet
 
 
 def _to_c_str_array(str_list):
@@ -40,7 +41,7 @@ def destroy_fen_batch(fen_batch):
 
 
 def create_sparse_batch_stream(
-    feature_set,
+    feature_set: str,
     concurrency,
     filenames,
     batch_size,
@@ -62,7 +63,7 @@ def destroy_sparse_batch_stream(stream):
     c_lib.dll.destroy_sparse_batch_stream(stream)
 
 
-def get_sparse_batch_from_fens(feature_set, fens, scores, plies, results):
+def get_sparse_batch_from_fens(feature_set: FeatureSet, fens, scores, plies, results):
     assert len(fens) == len(scores) == len(plies) == len(results)
 
     def to_c_int_array(data):
diff --git a/feature_transformer.py b/feature_transformer.py
@@ -573,13 +573,14 @@ def backward(ctx, grad_output_0, grad_output_1):
         return None, None, None, None, weight_grad, bias_grad
 
 
-class FeatureTransformerSlice(nn.Module):
+class BaseFeatureTransformerSlice(nn.Module):
     def __init__(self, num_inputs, num_outputs):
-        super(FeatureTransformerSlice, self).__init__()
+        super(BaseFeatureTransformerSlice, self).__init__()
         self.num_inputs = num_inputs
         self.num_outputs = num_outputs
 
         sigma = math.sqrt(1 / num_inputs)
+
         self.weight = nn.Parameter(
             torch.rand(num_inputs, num_outputs, dtype=torch.float32) * (2 * sigma)
             - sigma
@@ -588,27 +589,15 @@ def __init__(self, num_inputs, num_outputs):
             torch.rand(num_outputs, dtype=torch.float32) * (2 * sigma) - sigma
         )
 
+
+class FeatureTransformerSlice(BaseFeatureTransformerSlice):
     def forward(self, feature_indices, feature_values):
         return FeatureTransformerSliceFunction.apply(
             feature_indices, feature_values, self.weight, self.bias
         )
 
 
-class DoubleFeatureTransformerSlice(nn.Module):
-    def __init__(self, num_inputs, num_outputs):
-        super(DoubleFeatureTransformerSlice, self).__init__()
-        self.num_inputs = num_inputs
-        self.num_outputs = num_outputs
-
-        sigma = math.sqrt(1 / num_inputs)
-        self.weight = nn.Parameter(
-            torch.rand(num_inputs, num_outputs, dtype=torch.float32) * (2 * sigma)
-            - sigma
-        )
-        self.bias = nn.Parameter(
-            torch.rand(num_outputs, dtype=torch.float32) * (2 * sigma) - sigma
-        )
-
+class DoubleFeatureTransformerSlice(BaseFeatureTransformerSlice):
     def forward(
         self, feature_indices_0, feature_values_0, feature_indices_1, feature_values_1
     ):
@@ -624,8 +613,6 @@ def forward(
 
 if __name__ == "__main__":
     import time
-    import sys
-    import os
 
     def FeatureTransformerSliceFunctionEmulate(
         feature_indices, feature_values, weight, bias
diff --git a/ftperm.py b/ftperm.py
@@ -44,6 +44,8 @@
 import data_loader
 import model as M
 from model import NNUE
+from features.feature_set import FeatureSet
+
 
 """
 
@@ -394,7 +396,7 @@ def find_perm_impl(actmat, use_cupy):
 # -------------------------------------------------------------
 
 
-def read_model(nnue_path, feature_set):
+def read_model(nnue_path, feature_set: FeatureSet):
     with open(nnue_path, "rb") as f:
         reader = serialize.NNUEReader(f, feature_set)
         return reader.model
diff --git a/model.py b/model.py
@@ -1,9 +1,14 @@
 import ranger21
 import torch
-from torch import nn
+from torch import nn, Tensor
 import pytorch_lightning as pl
-from feature_transformer import DoubleFeatureTransformerSlice
+from feature_transformer import (
+    DoubleFeatureTransformerSlice,
+    BaseFeatureTransformerSlice,
+)
 from dataclasses import dataclass
+from features.feature_set import FeatureSet
+from typing import List, Tuple
 
 # 3 layer fully connected network
 L1 = 3072
@@ -24,25 +29,8 @@ class LossParams:
     qp_asymmetry: float = 0.0
 
 
-def coalesce_ft_weights(model, layer):
-    weight = layer.weight.data
-    indices = model.feature_set.get_virtual_to_real_features_gather_indices()
-    weight_coalesced = weight.new_zeros(
-        (model.feature_set.num_real_features, weight.shape[1])
-    )
-    for i_real, is_virtual in enumerate(indices):
-        weight_coalesced[i_real, :] = sum(
-            weight[i_virtual, :] for i_virtual in is_virtual
-        )
-    return weight_coalesced
-
-
-def get_parameters(layers):
-    return [p for layer in layers for p in layer.parameters()]
-
-
 class LayerStacks(nn.Module):
-    def __init__(self, count):
+    def __init__(self, count: int):
         super(LayerStacks, self).__init__()
 
         self.count = count
@@ -94,7 +82,7 @@ def _init_layers(self):
         self.output.weight = nn.Parameter(output_weight)
         self.output.bias = nn.Parameter(output_bias)
 
-    def forward(self, x, ls_indices):
+    def forward(self, x: Tensor, ls_indices: Tensor):
         assert self.idx_offset is not None and self.idx_offset.shape[0] == x.shape[0]
 
         indices = ls_indices.flatten() + self.idx_offset
@@ -162,7 +150,7 @@ class NNUE(pl.LightningModule):
 
     def __init__(
         self,
-        feature_set,
+        feature_set: FeatureSet,
         max_epoch=800,
         num_batches_per_epoch=int(100_000_000 / 16384),
         gamma=0.992,
@@ -304,7 +292,7 @@ def _clip_weights(self):
   to new_feature_set. Currently only works for adding virtual features.
   """
 
-    def set_feature_set(self, new_feature_set):
+    def set_feature_set(self, new_feature_set: FeatureSet):
         if self.feature_set.name == new_feature_set.name:
             return
 
@@ -351,14 +339,14 @@ def set_feature_set(self, new_feature_set):
 
     def forward(
         self,
-        us,
-        them,
-        white_indices,
-        white_values,
-        black_indices,
-        black_values,
-        psqt_indices,
-        layer_stack_indices,
+        us: Tensor,
+        them: Tensor,
+        white_indices: Tensor,
+        white_values: Tensor,
+        black_indices: Tensor,
+        black_values: Tensor,
+        psqt_indices: Tensor,
+        layer_stack_indices: Tensor,
     ):
         wp, bp = self.input(white_indices, white_values, black_indices, black_values)
         w, wpsqt = torch.split(wp, L1, dim=1)
@@ -382,7 +370,7 @@ def forward(
 
         return x
 
-    def step_(self, batch, batch_idx, loss_type):
+    def step_(self, batch: Tuple[Tensor, ...], batch_idx, loss_type):
         _ = batch_idx  # unused, but required by pytorch-lightning
 
         # We clip weights at the start of each step. This means that after
@@ -489,3 +477,20 @@ def configure_optimizers(self):
             optimizer, step_size=1, gamma=self.gamma
         )
         return [optimizer], [scheduler]
+
+
+def coalesce_ft_weights(model: NNUE, layer: BaseFeatureTransformerSlice):
+    weight = layer.weight.data
+    indices = model.feature_set.get_virtual_to_real_features_gather_indices()
+    weight_coalesced = weight.new_zeros(
+        (model.feature_set.num_real_features, weight.shape[1])
+    )
+    for i_real, is_virtual in enumerate(indices):
+        weight_coalesced[i_real, :] = sum(
+            weight[i_virtual, :] for i_virtual in is_virtual
+        )
+    return weight_coalesced
+
+
+def get_parameters(layers: List[nn.Module]):
+    return [p for layer in layers for p in layer.parameters()]
diff --git a/ranger21.py b/ranger21.py
diff --git a/serialize.py b/serialize.py
diff --git a/train.py b/train.py
