Implemented batching for TicTacToe AI

TicTacToe_AI/Net/pytorch_ai.py

@@ -4,6 +4,11 @@ import torch.optim as optim
 from torch import nn
 import torch.nn.functional as F
 from tqdm import tqdm
+import wandb
+
+wandb.init(project="tictactoe")
+
+BATCH_SIZE = 3
 
 
 def to_set(raw_list):
@@ -35,6 +40,40 @@ def to_set(raw_list):
     return out_set
 
 
+def to_batched_set(raw_list):
+    counter = 0
+    out_set = []
+    boardtensor = torch.zeros((BATCH_SIZE, 9))
+    labeltensor = torch.zeros(BATCH_SIZE, dtype=torch.long)
+    for line in tqdm(raw_list):
+        line = line.replace('\n', '')
+        raw_board, raw_label = line.split('|')[0], line.split('|')[1]
+
+        if not (int(raw_label) is -1):
+            labeltensor[counter] = int(raw_label)
+        else:
+            labeltensor[counter] = 9
+
+        raw_board = raw_board.split(',')
+        for n, block in enumerate(raw_board):
+            if int(block) is -1:
+                boardtensor[counter][n] = 0
+            elif int(block) is 0:
+                boardtensor[counter][n] = 0.5
+            elif int(block) is 1:
+                boardtensor[counter][n] = 1
+
+        if counter == (BATCH_SIZE - 1):
+            out_set.append([boardtensor, labeltensor])
+            boardtensor = torch.zeros((BATCH_SIZE, 9))
+            labeltensor = torch.zeros(BATCH_SIZE, dtype=torch.long)
+            counter = 0
+        else:
+            counter += 1
+
+    return out_set
+
+
 def buildsets():
     with open('boards.bds', 'r') as infile:
         print('Loading file...')
@@ -43,10 +82,10 @@ def buildsets():
         random.shuffle(alllines)
 
         print('Generating testset...')
-        testset = to_set(alllines[0:10000])
+        testset = to_batched_set(alllines[0:10000])
 
         print('Generating trainset...')
-        trainset = to_set(alllines[10001:200000])
+        trainset = to_batched_set(alllines[10001:20000])
 
     return trainset, testset
 
@@ -60,6 +99,7 @@ def testnet(net, testset):
             if torch.argmax(output) == label[0]:
                 correct += 1
             total += 1
+    wandb.log({'test_accuracy': correct / total})
     print("Accuracy: ", round(correct / total, 3))
 
 
@@ -79,21 +119,33 @@ class Net(torch.nn.Module):
         return F.log_softmax(x, dim=1)
 
 
-net = torch.load('./nets/net_3.pt')
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print('running on %s' % device)
+
+# net = torch.load('./nets/net_3.pt')
+
+net = Net()
+wandb.watch(net)
+
+net.to(device)
 
 optimizer = optim.Adam(net.parameters(), lr=0.001)
+loss_function = nn.CrossEntropyLoss()
 
 trainset, testset = buildsets()
 
 for epoch in range(100):
     print('Epoch: ' + str(epoch))
+    wandb.log({'epoch': epoch})
     for X, label in tqdm(trainset):
         net.zero_grad()
+        X.to(device)
         output = net(X)
-        loss = F.nll_loss(output.view(1, 10), label[0])
+        output.cpu()
+        loss = loss_function(output.view(-1, 10), label)
         loss.backward()
         optimizer.step()
+        wandb.log({'loss': loss})
 
-    print(loss)
-    torch.save(net, './nets/net_' + str(epoch + 3) + '.pt')
+    torch.save(net, './nets/gpunets/net_' + str(epoch) + '.pt')
     testnet(net, testset)