Increasing WildBgDataSet speed (#10)

Moving away from pandas dataframe changing to loading data into tensor
on init

training/src/dataset.py

@@ -0,0 +1,27 @@
+from torch.utils.data import Dataset
+import torch
+
+class WildBgDataSet(Dataset):
+    def __init__(self, csv_files: list | str):
+        # If you want to combine several CSV files from another folder, use the following: 
+        if isinstance(csv_files, str):
+            csv_files = [csv_files]
+        labels = []
+        inputs = []
+        for path in csv_files:
+            with open(path, 'r') as f:
+                lines = f.readlines()
+                for line in lines[1:]:
+                    line = line.strip().split(';')
+                    line = list(map(float, line))
+                    labels.append(line[:6])
+                    inputs.append(line[6:])
+        self.inputs = torch.Tensor(inputs)
+        self.labels = torch.Tensor(labels)
+
+    def __len__(self):
+        return self.inputs.shape[0]
+
+    def __getitem__(self, idx):
+        # First 6 columns are outputs, last 202 columns are inputs
+        return self.inputs[idx], self.labels[idx]

training/src/model.py

@@ -0,0 +1,29 @@
+from torch import nn
+
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        # Inputs to hidden layer linear transformation
+        self.hidden1 = nn.Linear(202, 300)
+        self.hidden2 = nn.Linear(300, 250)
+        self.hidden3 = nn.Linear(250, 200)
+
+        # Output layer, 6 outputs for win/lose - normal/gammon/bg
+        self.output = nn.Linear(200, 6)
+
+        # Define activation function and softmax output 
+        self.activation = nn.ReLU()
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        # Pass the input tensor through each of our operations
+        x = self.hidden1(x)
+        x = self.activation(x)
+        x = self.hidden2(x)
+        x = self.activation(x)
+        x = self.hidden3(x)
+        x = self.activation(x)
+        x = self.output(x)
+        x = self.softmax(x)
+        return x

training/src/train-on-rollout-data.py

@@ -2,98 +2,62 @@
 import torch
 from torch import nn
 from torch.utils.data import DataLoader
-from torch.utils.data import Dataset
-import pandas as pd
-
-# "mps" takes more time than "cpu" on Macs, so let's ignore it for now.
-device = (
-    "cuda"
-    if torch.cuda.is_available()
-    # else "mps"
-    # if torch.backends.mps.is_available()
-    else "cpu"
-)
-print(f"Using {device} device")
-
-class WildBgDataSet(Dataset):
-    def __init__(self):
-        # If you want to combine several CSV files from another folder, use the following: 
-        # csv_files = [
-        #     "../../wildbg-training/data/0006/rollouts.csv",
-        #     "../../wildbg-training/data/0007/rollouts.csv",
-        # ]
-        csv_files = ["../training-data/rollouts.csv"]
-        self.data = pd.concat([pd.read_csv(f, sep=';') for f in csv_files ], ignore_index=True)
-
-    def __len__(self):
-        return self.data.shape[0]
-
-    def __getitem__(self, idx):
-        # First 6 columns are outputs, last 202 columns are inputs
-        output = self.data.iloc[idx, 0:6]
-        input = self.data.iloc[idx, 6:208]
-        return torch.FloatTensor(input).to(device), torch.FloatTensor(output).to(device)
-
-
-class Network(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        # Inputs to hidden layer linear transformation
-        self.hidden1 = nn.Linear(202, 300)
-        self.hidden2 = nn.Linear(300, 250)
-        self.hidden3 = nn.Linear(250, 200)
-
-        # Output layer, 6 outputs for win/lose - normal/gammon/bg
-        self.output = nn.Linear(200, 6)
+from model import Model
+from dataset import WildBgDataSet
+
+def save_model(model: nn.Module, path: str) -> None:
+    dummy_input = torch.randn(1, 202, requires_grad=True, device=device)
+    torch.onnx.export(model, dummy_input, path)
+
+def train(model: nn.Module, trainloader: DataLoader, epochs: int) -> nn.Module:
+    # Define loss function, L1Loss and MSELoss are good choices
+    criterion = nn.MSELoss()
+
+    # Optimizer based on model, adjust the learning rate
+    # 4.0 has worked well for Tanh(), one layer and 100k positions
+    # 3.0 has worked well for ReLu(), three layers and 200k positions
+    optimizer = torch.optim.SGD(model.parameters(), lr=3.0)
+
+    for epoch in range(epochs):
+        epoch_loss = 0.0
+        for i, data in enumerate(trainloader, 0):
+            inputs, labels = data
+            # set optimizer to zero grad to remove previous epoch gradients
+            optimizer.zero_grad()
+            # forward propagation
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            # backward propagation
+            loss.backward()
+            # optimize
+            optimizer.step()
+            epoch_loss += loss.item()
 
-        # Define activation function and softmax output 
-        self.activation = nn.ReLU()
-        self.softmax = nn.Softmax(dim=1)
-
-    def forward(self, x):
-        # Pass the input tensor through each of our operations
-        x = self.hidden1(x)
-        x = self.activation(x)
-        x = self.hidden2(x)
-        x = self.activation(x)
-        x = self.hidden3(x)
-        x = self.activation(x)
-        x = self.output(x)
-        x = self.softmax(x)
-        return x
-
-model = Network().to(device)
-
-traindata = WildBgDataSet()
-trainloader = DataLoader(traindata, batch_size=64, shuffle=True)
-
-# Define loss function, L1Loss and MSELoss are good choices
-criterion = nn.MSELoss()
-
-# Optimizer based on model, adjust the learning rate
-# 4.0 has worked well for Tanh(), one layer and 100k positions
-# 3.0 has worked well for ReLu(), three layers and 200k positions
-optimizer = torch.optim.SGD(model.parameters(), lr=3.0)
-
-epochs = 20
-
-for epoch in range(epochs):
-    running_loss = 0.0
-    for i, data in enumerate(trainloader, 0):
-        inputs, labels = data
-        # set optimizer to zero grad to remove previous epoch gradients
-        optimizer.zero_grad()
-        # forward propagation
-        outputs = model(inputs)
-        loss = criterion(outputs, labels)
-        # backward propagation
-        loss.backward()
-        # optimize
-        optimizer.step()
-        running_loss += loss.item()
-        print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.5f}')
-
-Path("../neural-nets").mkdir(exist_ok=True)
-dummy_input = torch.randn(1, 202, requires_grad=True, device=device)
-model_onnx = torch.onnx.export(model, dummy_input, "../neural-nets/wildbg.onnx")
+        epoch_loss /= len(trainloader) / 64
+        print(f'[Epoch: {epoch + 1}] loss: {epoch_loss:.5f}')
+
+    return model
+
+def main(model: nn.Module, data_path: str, model_path: str):
+    traindata = WildBgDataSet(data_path)
+    trainloader = DataLoader(traindata, batch_size=64, shuffle=True)
+
+    try:
+        model = train(model, trainloader, 20)
+    finally:
+        print('Finished Training')
+        save_model(model, model_path)
+
+if __name__ == "__main__":
+    # "mps" takes more time than "cpu" on Macs, so let's ignore it for now.
+    device = (
+        "cuda"
+        if torch.cuda.is_available()
+        # else "mps"
+        # if torch.backends.mps.is_available()
+        else "cpu"
+    )
+    print(f"Using {device} device")
+    Path("../neural-nets").mkdir(exist_ok=True)
+    model = Model().to(device)
+    main(model, "../training-data/rollouts.csv", "../neural-nets/wildbg.onnx")