New NNUE architecture, (768->1024)x2->1

bmdanielsson · bmdanielsson · commit ac8d6e108e7d · 2024-05-15T18:20:56.000+02:00
diff --git a/model.py b/model.py
@@ -1,62 +1,38 @@
-import torch
-import struct
-
-from torch import nn
-
-
-# Number of inputs
-NUM_SQ = 64
-NUM_PT = 12
-NUM_INPUTS = NUM_SQ*NUM_PT
-
-# 3 layer fully connected network
-L1 = 384
-L2 = 8
-L3 = 16
-
-class NNUE(nn.Module):
-    def __init__(self):
-        super(NNUE, self).__init__()
-        self.input = nn.Linear(NUM_INPUTS, L1)
-        self.l1 = nn.Linear(2 * L1, L2)
-        self.l2 = nn.Linear(L2, L3)
-        self.output = nn.Linear(L3, 1)
-
-
-    def forward(self, us, them, w_in, b_in):
-        w = self.input(w_in)
-        b = self.input(b_in)
-        l0_ = (us*torch.cat([w, b], dim=1)) + (them*torch.cat([b, w], dim=1))
-        l0_ = torch.clamp(l0_, 0.0, 1.0)
-        l1_ = torch.clamp(self.l1(l0_), 0.0, 1.0)
-        l2_ = torch.clamp(self.l2(l1_), 0.0, 1.0)
-        x = self.output(l2_)
-        return x
-
-
-    def clamp_weights(self):
-        # L1
-        data = self.l1.weight.data
-        data.clamp(-127.0/64.0, 127.0/64.0)
-        self.l1.weight.data.copy_(data)
-
-        # L2
-        data = self.l2.weight.data
-        data.clamp(-127.0/64.0, 127.0/64.0)
-        self.l2.weight.data.copy_(data)
-
-        # Output
-        data = self.output.weight.data
-        data.clamp(-127.0*127.0/64.0, 127.0*127.0/64.0)
-        self.output.weight.data.copy_(data)
-
-
-def loss_function(wdl, pred, batch):
-    us, them, white, black, outcome, score = batch
-    
-    wdl_eval_model = (pred*600.0/361).sigmoid()
-    wdl_eval_target = (score/410).sigmoid()
-
-    wdl_value_target = wdl_eval_target * (1.0 - wdl) + outcome * wdl
-    
-    return torch.abs(wdl_value_target  - wdl_eval_model).square().mean()
+import torch
+import struct
+
+from torch import nn
+
+
+# Number of inputs
+NUM_SQ = 64
+NUM_PT = 12
+NUM_INPUTS = NUM_SQ*NUM_PT
+
+L1 = 1024
+
+class NNUE(nn.Module):
+    def __init__(self):
+        super(NNUE, self).__init__()
+        self.input = nn.Linear(NUM_INPUTS, L1)
+        self.output = nn.Linear(2*L1, 1)
+
+
+    def forward(self, us, them, w_in, b_in):
+        w = self.input(w_in)
+        b = self.input(b_in)
+        l0_ = (us*torch.cat([w, b], dim=1)) + (them*torch.cat([b, w], dim=1))
+        l0_ = torch.clamp(l0_, 0.0, 1.0)
+        x = self.output(l0_)
+        return x
+
+
+def loss_function(wdl, pred, batch):
+    us, them, white, black, outcome, score = batch
+    
+    wdl_eval_model = (pred*600.0/361).sigmoid()
+    wdl_eval_target = (score/410).sigmoid()
+
+    wdl_value_target = wdl_eval_target * (1.0 - wdl) + outcome * wdl
+    
+    return torch.abs(wdl_value_target  - wdl_eval_model).square().mean()
diff --git a/quantize.py b/quantize.py
@@ -12,14 +12,11 @@
 
 HALFKX_WEIGHT_SCALE = MAX_QUANTIZED_ACTIVATION
 HALFKX_BIAS_SCALE = MAX_QUANTIZED_ACTIVATION
-HIDDEN_WEIGHT_SCALE = (1<<WEIGHT_SCALE_BITS)
-HIDDEN_BIAS_SCALE = (1<<WEIGHT_SCALE_BITS)*MAX_QUANTIZED_ACTIVATION
 OUTPUT_WEIGHT_SCALE = (OUTPUT_SCALE*NNUE2SCORE/MAX_QUANTIZED_ACTIVATION)
 OUTPUT_BIAS_SCALE = OUTPUT_SCALE*NNUE2SCORE
-MAX_HIDDEN_WEIGHT = MAX_QUANTIZED_ACTIVATION/HIDDEN_WEIGHT_SCALE
 MAX_OUTPUT_WEIGHT = MAX_QUANTIZED_ACTIVATION/OUTPUT_WEIGHT_SCALE
 
-NNUE_FORMAT_VERSION = 0x00000009
+NNUE_FORMAT_VERSION = 0x0000000A
 
 def write_header(buf, version):
     buf.extend(struct.pack('<I', version))
@@ -41,12 +38,6 @@ def quant_input(biases, weights):
     return (biases, weights)
 
 
-def quant_linear(biases, weights):
-    biases = biases.data.mul(HIDDEN_BIAS_SCALE).round().to(torch.int32)
-    weights = weights.data.clamp(-MAX_HIDDEN_WEIGHT, MAX_HIDDEN_WEIGHT).mul(HIDDEN_WEIGHT_SCALE).round().to(torch.int8)
-    return (biases, weights)
-
-
 def quant_output(biases, weights):
     biases = biases.data.mul(OUTPUT_BIAS_SCALE).round().to(torch.int32)
     weights = weights.data.clamp(-MAX_OUTPUT_WEIGHT, MAX_OUTPUT_WEIGHT).mul(OUTPUT_WEIGHT_SCALE).round().to(torch.int8)
@@ -63,16 +54,12 @@ def quantization(source, target):
 
     # Perform quantization
     input = quant_input(nnue.input.bias, nnue.input.weight)
-    linear1 = quant_linear(nnue.l1.bias, nnue.l1.weight)
-    linear2 = quant_linear(nnue.l2.bias, nnue.l2.weight)
     output = quant_output(nnue.output.bias, nnue.output.weight)
 
     # Write quantized layers
     outbuffer = bytearray()
     write_header(outbuffer, NNUE_FORMAT_VERSION)
     write_input(outbuffer, input[0], input[1])
-    write_layer(outbuffer, linear1[0], linear1[1])
-    write_layer(outbuffer, linear2[0], linear2[1])
     write_layer(outbuffer, output[0], output[1])
     with open(target, 'wb') as f:
         f.write(outbuffer)
