Change format for saved snapshots

bmdanielsson · bmdanielsson · commit ae5aab4b51c6 · 2024-01-28T17:52:10.000+01:00
Change snapshot format to Pytorch state dictionary.
diff --git a/model.py b/model.py
@@ -3,9 +3,6 @@
 from torch import nn
 import torch.nn.functional as F
 
-# The version of the export format
-EXPORT_FORMAT_VERSION = 0x00000008
-
 # Number of inputs
 NUM_SQ = 64
 NUM_PT = 12
@@ -36,31 +33,6 @@ def forward(self, us, them, w_in, b_in):
         return x
 
 
-    def serialize_halfkx_layer(self, buf, layer):
-        bias = layer.bias.data.cpu()
-        buf.extend(bias.flatten().numpy().tobytes())
-        weight = self.input.weight.data.clone().cpu()
-        buf.extend(weight.transpose(0, 1).flatten().numpy().tobytes())
-
-
-    def serialize_linear_layer(self, buf, layer):
-        bias = layer.bias.data.cpu()
-        buf.extend(bias.flatten().numpy().tobytes())
-        weight = layer.weight.data.cpu()
-        buf.extend(weight.flatten().numpy().tobytes())
-
-
-    def serialize(self, buf):
-        # Write header
-        buf.extend(struct.pack('<i', EXPORT_FORMAT_VERSION))
-
-        # Write layers
-        self.serialize_halfkx_layer(buf, self.input)
-        self.serialize_linear_layer(buf, self.l1)
-        self.serialize_linear_layer(buf, self.l2)
-        self.serialize_linear_layer(buf, self.output)
-
-
 def loss_function(wdl, pred, batch):
     us, them, white, black, outcome, score = batch
     
diff --git a/nettest.py b/nettest.py
@@ -71,7 +71,7 @@ def main(args):
         # Find all .bin files in the net folder
         bin_files = [f for f in os.listdir(args.net_dir)
                        if (os.path.isfile(os.path.join(args.net_dir, f)) and
-                           os.path.splitext(f)[1] == '.bin')]
+                           os.path.splitext(f)[1] == '.pt')]
         time.sleep(5)
 
         # Run a match with each new net
diff --git a/quantize.py b/quantize.py
@@ -19,6 +19,7 @@
 MAX_HIDDEN_WEIGHT = MAX_QUANTIZED_ACTIVATION/HIDDEN_WEIGHT_SCALE
 MAX_OUTPUT_WEIGHT = MAX_QUANTIZED_ACTIVATION/OUTPUT_WEIGHT_SCALE
 
+NNUE_FORMAT_VERSION = 0x00000008
 
 def write_header(buf, version):
     buf.extend(struct.pack('<I', version))
@@ -29,7 +30,7 @@ def write_layer(buf, biases, weights):
     buf.extend(weights.numpy().tobytes())
 
 
-def quant_halfkx(biases, weights):
+def quant_input(biases, weights):
     biases = biases.mul(HALFKX_BIAS_SCALE).round().to(torch.int16)
     weights = weights.mul(HALFKX_WEIGHT_SCALE).round().to(torch.int16)
     return (biases, weights)
@@ -47,14 +48,7 @@ def quant_output(biases, weights):
     return (biases, weights)
 
 
-def read_version(file):
-    version = struct.unpack('<I', file.read(4))[0]
-    if version != model.EXPORT_FORMAT_VERSION:
-        raise Exception('Model format mismatch')
-    return version
-    
-
-def read_layer(file, ninputs, size):
+def extract_layer(file, ninputs, size):
     buf = numpy.fromfile(file, numpy.float32, size)
     biases = torch.from_numpy(buf.astype(numpy.float32))
     buf = numpy.fromfile(file, numpy.float32, size*ninputs)
@@ -65,24 +59,21 @@ def read_layer(file, ninputs, size):
 def quantization(source, target):
     print('Performing quantization ...')
 
-    # Read all layers
-    with open(source, 'rb') as f:
-        version = read_version(f)
-        halfkx = read_layer(f, model.NUM_INPUTS, model.L1)
-        linear1 = read_layer(f, model.L1*2, model.L2)
-        linear2 = read_layer(f, model.L2, model.L3)
-        output = read_layer(f, model.L3, 1)
+    # Load model
+    nnue = model.NNUE()
+    nnue.load_state_dict(torch.load(source, map_location=torch.device('cpu')))
+    nnue.eval()
 
     # Perform quantization
-    halfkx = quant_halfkx(halfkx[0], halfkx[1])
-    linear1 = quant_linear(linear1[0], linear1[1])
-    linear2 = quant_linear(linear2[0], linear2[1])
-    output = quant_output(output[0], output[1])
+    input = quant_input(nnue.input.weight, nnue.input.bias)
+    linear1 = quant_linear(nnue.l1.weight, nnue.l1.bias)
+    linear2 = quant_linear(nnue.l2.weight, nnue.l2.bias)
+    output = quant_output(nnue.output.weight, nnue.output.bias)
 
     # Write quantized layers
     outbuffer = bytearray()
-    write_header(outbuffer, version)
-    write_layer(outbuffer, halfkx[0], halfkx[1])
+    write_header(outbuffer, NNUE_FORMAT_VERSION)
+    write_layer(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])
diff --git a/train.py b/train.py
@@ -29,18 +29,12 @@ def filter_saved_models(saved_models, top_n):
 
 
 def write_model(nnue, path):
-    # Serialize the model to a buffer
-    buf = bytearray()
-    nnue.serialize(buf)
-
-    # Write the buffer
-    with open(path, 'wb') as f:
-        f.write(buf)
+    torch.save(nnue.state_dict(), path)
 
 
 def save_model(nnue, output_path, epoch, idx, val_loss):
     # Construct the full path
-    path = f'{output_path}/epoch_{epoch}_iter_{idx+1}_loss_{val_loss:.5f}.bin.tmp'
+    path = f'{output_path}/epoch_{epoch}_iter_{idx+1}_loss_{val_loss:.5f}.pt.tmp'
 
     # Save the model
     write_model(nnue, path)
