train

train/README

@@ -0,0 +1,11 @@
+# Train
+
+## Dependencies
+
+python 3.13
+pytorch 2.6.0+cu126
+libzmq
+
+## Run
+
+python train.py --game go9

train/dataloader.py

@@ -0,0 +1,85 @@
+from contextlib import contextmanager
+import os
+import re
+from time import time
+import torch
+from rb import ReplayBuffer, SocketManager
+
+
+@contextmanager
+def timer(header):
+    time_start = time()
+    yield
+    print(header, (time() - time_start))
+
+
+class DataLoader:
+    def __init__(self, port, cpus, batch_size, sgf_prefix, is_selfplay):
+        seed = torch.randint(0, 2**32, [1]).item()
+        rb = ReplayBuffer()
+        rb.max_iters = 20
+        # moves_per_iter = #games * #moves_per_game
+        rb.moves_per_iter = 5000 * 60
+        rb.run(seed, cpus, batch_size)
+        sock = SocketManager()
+        sock.run(port)
+        self._rb = rb
+        self._sock = sock
+        # total = #moves_per_game * #isom / batch_size
+        self._total = int(rb.moves_per_iter * 1 / batch_size)
+        self._sgf_prefix = sgf_prefix
+        self._iter = 0
+        self._is_selfplay = is_selfplay
+        if not is_selfplay:
+            import zmq
+
+            self._ctx = zmq.Context()
+            zsock = self._ctx.socket(zmq.DEALER)
+            zsock.setsockopt(zmq.LINGER, 0)
+            zsock.setsockopt(zmq.ROUTING_ID, b"0")
+            zsock.connect(f"tcp://127.0.0.1:{port}")
+            zsock.send_multipart([b"0", b"0"])
+            self._zsock = zsock
+
+    def load(self, sgf_prefix, epoch_ckpt):
+        rb, sock = self._rb, self._sock
+        for i in range(epoch_ckpt):
+            if self._iter > 0:
+                sock.notify()
+            if i == 0 or i + rb.max_iters >= epoch_ckpt:
+                print(f"[{i:3d}] Load selfplay")
+                pattern = re.compile(rf"iter-{self._iter}-(\d+).sgf")
+                nodes = 1 + max(
+                    int(m.group(1))
+                    for f in os.listdir(sgf_prefix)
+                    if (m := pattern.search(f))
+                )
+                rb.add_iter(sgf_prefix, self._iter, nodes)
+                for _ in range(self._total):
+                    rb.sample().free()
+            self._iter += 1
+
+    def __del__(self):
+        self._rb.terminate()
+        self._sock.terminate()
+
+    def __iter__(self):
+        rb, sock = self._rb, self._sock
+        if self._iter > 0:
+            sock.notify()
+        if self._is_selfplay:
+            with timer("[{:3d}] Time for selfplay:".format(self._iter)):
+                if sock.wait():
+                    exit(0)  # SIGINT
+        else:
+            finished = rb.moves_per_iter + 1
+            self._zsock.send_multipart(
+                [bytes(str(self._iter), "utf-8"), bytes(str(finished), "utf-8")]
+            )
+        with timer("[{:3d}] Time for training:".format(self._iter)):
+            rb.add_iter(self._sgf_prefix, self._iter, sock.nodes)
+            for _ in range(self._total):
+                sample = rb.sample()
+                yield sample
+                sample.free()
+        self._iter += 1

train/model.py

@@ -0,0 +1,165 @@
+from torch import nn
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channels, channels, bias, k=3, p=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, channels, k, stride=1, padding=p, bias=bias)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv2d(channels, channels, k, stride=1, padding=p, bias=bias)
+        self.bn2 = nn.BatchNorm2d(channels)
+        self.relu2 = nn.ReLU()
+
+    def forward(self, x):
+        y = self.conv1(x)
+        y = self.bn1(y)
+        y = self.relu1(y)
+        y = self.conv2(y)
+        y = self.bn2(y)
+        x = x + y
+        x = self.relu2(x)
+        return x
+
+
+class Bottleneck(nn.Module):
+    def __init__(self, in_channels, channels, bias):
+        super().__init__()
+        mid_channels = channels // 2
+        self.conv1 = nn.Conv2d(in_channels, mid_channels, 1, 1, bias=bias)
+        self.bn1 = nn.BatchNorm2d(mid_channels)
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv2d(mid_channels, mid_channels, 3, 1, padding=1, bias=bias)
+        self.bn2 = nn.BatchNorm2d(mid_channels)
+        self.relu2 = nn.ReLU()
+        self.conv3 = nn.Conv2d(mid_channels, channels, 1, 1, bias=bias)
+        self.bn3 = nn.BatchNorm2d(channels)
+        self.relu3 = nn.ReLU()
+
+    def forward(self, x):
+        y = self.conv1(x)
+        y = self.bn1(y)
+        y = self.relu1(y)
+        y = self.conv2(y)
+        y = self.bn2(y)
+        y = self.relu2(y)
+        y = self.conv3(y)
+        y = self.bn3(y)
+        x = x + y
+        x = self.relu3(x)
+        return x
+
+
+class Bottlenest(nn.Module):
+    def __init__(self, in_channels, channels, bias):
+        super().__init__()
+        mid_channels = channels // 2
+        self.conv0 = nn.Conv2d(in_channels, mid_channels, 1, 1, bias=bias)
+        self.bn0 = nn.BatchNorm2d(mid_channels)
+        self.conv1 = nn.Conv2d(mid_channels, mid_channels, 3, 1, padding=1, bias=bias)
+        self.bn1 = nn.BatchNorm2d(mid_channels)
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv2d(mid_channels, mid_channels, 3, 1, padding=1, bias=bias)
+        self.bn2 = nn.BatchNorm2d(mid_channels)
+        self.relu2 = nn.ReLU()
+        self.conv3 = nn.Conv2d(mid_channels, mid_channels, 3, 1, padding=1, bias=bias)
+        self.bn3 = nn.BatchNorm2d(mid_channels)
+        self.relu3 = nn.ReLU()
+        self.conv4 = nn.Conv2d(mid_channels, mid_channels, 3, 1, padding=1, bias=bias)
+        self.bn4 = nn.BatchNorm2d(mid_channels)
+        self.relu4 = nn.ReLU()
+        self.conv5 = nn.Conv2d(mid_channels, channels, 1, 1, bias=bias)
+        self.bn5 = nn.BatchNorm2d(channels)
+        self.relu5 = nn.ReLU()
+
+    def forward(self, x):
+        y = self.conv0(x)
+        y = self.bn0(y)
+        z = self.conv1(y)
+        z = self.bn1(z)
+        z = self.relu1(z)
+        z = self.conv2(z)
+        z = self.bn2(z)
+        y = y + z
+        y = self.relu2(y)
+        z = self.conv3(y)
+        z = self.bn3(z)
+        z = self.relu3(z)
+        z = self.conv4(z)
+        z = self.bn4(z)
+        y = y + z
+        y = self.relu4(y)
+        y = self.conv5(y)
+        y = self.bn5(y)
+        x = x + y
+        x = self.relu5(x)
+        return x
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, in_channels, layers, channels, bias):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels, channels, kernel_size=5, stride=1, padding=2, bias=bias
+            ),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+        )
+        self.convs = nn.ModuleList(
+            [block(channels, channels, bias) for _ in range(layers)]
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        for conv in self.convs:
+            x = conv(x)
+        return x
+
+
+class AlphaZero(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        layers,
+        channels,
+        moves,
+        board_size,
+        value_heads=1,
+        bias=False,
+        block=BasicBlock,
+    ):
+        super().__init__()
+        self.board_size = board_size
+        self.resnet = ResNet(block, in_channels, layers, channels, bias)
+        # policy head
+        self.policy_head_front = nn.Sequential(
+            nn.Conv2d(channels, 2, 1),
+            nn.BatchNorm2d(2),
+            nn.ReLU(),
+        )
+        self.policy_head_end = nn.Linear(2 * board_size, moves)
+        # value head
+        self.value_head_front = nn.Sequential(
+            nn.Conv2d(channels, 1, 1),
+            nn.BatchNorm2d(1),
+            nn.ReLU(),
+        )
+        self.value_head_end = nn.Sequential(
+            nn.Linear(board_size, channels),
+            nn.ReLU(),
+            nn.Linear(channels, value_heads),
+            nn.Tanh(),
+        )
+
+    def forward(self, x):
+        x = self.resnet(x)
+        # policy head
+        p = self.policy_head_front(x)
+        p = p.view(-1, 2 * self.board_size)
+        p = self.policy_head_end(p)
+        # value head
+        v = self.value_head_front(x)
+        v = v.view(-1, self.board_size)
+        v = self.value_head_end(v)
+        return p, v

train/train.py

@@ -0,0 +1,123 @@
+from dataloader import DataLoader
+from model import AlphaZero, BasicBlock, Bottlenest
+#from export_ait import save_ait
+import argparse
+import os
+import re
+import time
+import torch
+from torch import nn
+
+kGames = dict(
+    nogo=dict(num_features=4, moves=81, board_size=81, value_heads=1),
+    go9=dict(num_features=20, moves=82, board_size=81, value_heads=31),
+    go19=dict(num_features=20, moves=362, board_size=361, value_heads=31),
+)
+
+
+def save_model(model_prefix, epoch, net, optimizer, moves, board_size):
+    net.eval()
+    net_state = net.state_dict()
+    torch.save(
+        {
+            "epoch": epoch,
+            "net": net_state,
+            "optimizer": optimizer.state_dict(),
+        },
+        f"{model_prefix}/model-{epoch}.ckpt",
+    )
+    #save_ait(net_state, moves, board_size, f"{model_prefix}/model-{epoch}.ait")
+    net.train()
+
+
+def main(args):
+    torch.backends.cudnn.benchmark = True
+    game = kGames[args.game]
+    moves, board_size = game["moves"], game["board_size"]
+    layers, channels, block = re.search(r"b(\d+)c(\d+)(.*)", args.model_prefix).groups()
+    block = BasicBlock if block == "" else Bottlenest
+    net = AlphaZero(
+        in_channels=game["num_features"],
+        layers=int(layers),
+        channels=int(channels),
+        moves=moves,
+        board_size=board_size,
+        value_heads=game["value_heads"],
+        bias=False,
+        block=block,
+    ).cuda()
+    # loss fn
+    p_criterion = lambda p_logits, p_labels: (
+        (-p_labels * torch.log_softmax(p_logits, dim=1)).sum(dim=1).mean()
+    )
+    v_criterion = nn.MSELoss()
+    optimizer = torch.optim.SGD(
+        net.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.0001, nesterov=True
+    )
+    # load checkpoint
+    epoch_start = 0
+    dataloader = DataLoader(
+        args.port, args.cpus, args.batch_size, args.sgf_prefix, not args.pretrain
+    )
+    if args.load_ckpt:
+        print("> Restore from", args.load_ckpt)
+        ckpt = torch.load(args.load_ckpt, weights_only=True)
+        net.load_state_dict(ckpt["net"])
+        optimizer.load_state_dict(ckpt["optimizer"])
+        if args.load_data:
+            epoch_start = ckpt["epoch"]
+            dataloader.load(args.load_data, epoch_start)
+    save_model(args.model_prefix, epoch_start, net, optimizer, moves, board_size)
+    print("> Start training")
+    # train
+    for epoch in range(epoch_start, epoch_start + 6000):
+        net.train()
+        time_start = time.time()
+        for i, batch in enumerate(dataloader):
+            inputs, p_labels, v_labels = batch.inputs, batch.policy, batch.value
+
+            # forward + backward
+            p_logits, v_logits = net(inputs)
+            
+            v_loss = v_criterion(v_logits, v_labels)
+            p_loss = p_criterion(p_logits, p_labels)
+            loss = v_loss * args.value_ratio + p_loss
+
+            # optimize
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # train loss
+            if i % 10 == 0:
+                print(
+                    "[{:3d}:{:5d}] PN_Loss: {:.5f} VN_Loss: {:.5f}".format(
+                        epoch, i, p_loss.item(), v_loss.item()
+                    )
+                )
+
+        print("[{:3d}] Time per epoch: {}".format(epoch, time.time() - time_start))
+        save_model(args.model_prefix, epoch + 1, net, optimizer, moves, board_size)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # game
+    parser.add_argument("--game", default="nogo")
+    # training
+    parser.add_argument("--pretrain", action="store_true")
+    parser.add_argument("--sgf-prefix", default="../selfplay/sp")
+    parser.add_argument("--model-prefix", default="models_b6c96")
+    parser.add_argument("--load-ckpt", default="")
+    parser.add_argument("--load-data", default="")
+    parser.add_argument("--cpus", default=32, type=int)
+    parser.add_argument("--port", default=5566, type=int)
+    # hyperparameters
+    parser.add_argument("-lr", "--lr", default=0.01, type=float)
+    parser.add_argument("-bs", "--batch-size", default=512, type=int)
+    parser.add_argument("-vr", "--value-ratio", default=1, type=float)
+
+    args = parser.parse_args()
+    os.makedirs(args.sgf_prefix, exist_ok=True)
+    os.makedirs(args.model_prefix, exist_ok=True)
+    main(args)