Simplify

martinnovaak · web-flow · commit cabec10d4f87 · 2024-02-03T16:44:23.000+01:00
diff --git a/trainer/batchloader.py b/trainer/batchloader.py
@@ -30,25 +30,17 @@ class Batch:
 class BatchLoader:
     def __init__(self, lib_path: str, files: list[bytes], batch_size: int, scale: float, wdl: float) -> None:
         self.parse_lib = None
-        if not files:
-            raise ValueError("The files list cannot be empty.")
-
-        try:
-            self.parse_lib = ctypes.CDLL(lib_path)
-        except OSError as e:
-            raise Exception(f"Failed to load the library: {e}")
+        if not files: raise ValueError("The files list cannot be empty.")
+        try: self.parse_lib = ctypes.CDLL(lib_path)
+        except OSError as e: raise Exception(f"Failed to load the library: {e}")
         self.load_parse_lib()
 
-        self.files = files
-        self.file_index = 0
+        self.files, self.file_index = files, 0
         self.batch = ctypes.c_void_p(self.parse_lib.batch_new(ctypes.c_uint32(batch_size), ctypes.c_float(scale), ctypes.c_float(wdl)))
-
-        if self.batch.value is None:
-            raise Exception("Failed to create batch")
+        if self.batch.value is None: raise Exception("Failed to create batch")
 
         self.current_reader = ctypes.c_void_p(self.parse_lib.file_reader_new(ctypes.create_string_buffer(files[0])))
-        if self.current_reader.value is None:
-            raise Exception("Failed to create file reader")
+        if self.current_reader.value is None: raise Exception("Failed to create file reader")
 
     def next_batch(self, device: torch.device) -> tuple[bool, Batch]:
         new_epoch = False
diff --git a/trainer/model.py b/trainer/model.py
@@ -24,22 +24,3 @@ def forward(self, batch: Batch):
     def clamp_weights(self):
         self.feature_transformer.weight.data.clamp_(-2.0, 2.0)
         self.output_layer.weight.data.clamp_(-2.0, 2.0)
-
-
-# 768 -> N -> 1
-class StmPerspectiveNetwork(torch.nn.Module):
-    def __init__(self, feature_output_size: int):
-        super().__init__()
-        self.feature_transformer = torch.nn.Linear(768, feature_output_size)
-        self.output_layer = torch.nn.Linear(feature_output_size, 1)
-
-    def forward(self, batch: Batch):
-        stm_perspective = self.feature_transformer(batch.stm_sparse.to_dense())
-
-        hidden_features = torch.clamp(stm_perspective, 0, 1)
-
-        return torch.sigmoid(self.output_layer(hidden_features))
-
-    def clamp_weights(self):
-        self.feature_transformer.weight.data.clamp_(-2.0, 2.0)
-        self.output_layer.weight.data.clamp_(-2.0, 2.0)
diff --git a/trainer/train.py b/trainer/train.py
@@ -1,8 +1,7 @@
 from time import time
-
 import torch
-from batchloader import BatchLoader
 
+from batchloader import BatchLoader
 from quantize import quantize
 
 
@@ -12,13 +11,7 @@ def print_epoch_stats(epoch, running_loss, iterations, fens, start_time, current
                .format(epoch, epoch_time, running_loss.item() / iterations, fens / epoch_time))
     print(message)
 
-def train(
-        model: torch.nn.Module,
-        optimizer: torch.optim.Optimizer,
-        dataloader: BatchLoader,
-        epochs: int,
-        device: torch.device,
-) -> None:
+def train(model: torch.nn.Module, optimizer: torch.optim.Optimizer, dataloader: BatchLoader, epochs: int, device: torch.device):
     running_loss = torch.zeros(1, device=device)
     epoch_start_time = time()
     iterations = 0
@@ -37,7 +30,7 @@ def train(
             epoch_start_time = current_time
             iterations = 0
             fens = 0
-            
+
             quantize(model, f"network/nnue_{epoch}_scaled.bin")
 
         optimizer.zero_grad()
