50 lines
1.3 KiB (Stored with Git LFS)
Python
50 lines
1.3 KiB (Stored with Git LFS)
Python
import torch
|
|
import torch.nn.functional as F
|
|
import torch.nn as nn
|
|
|
|
import numpy as np
|
|
|
|
TARGET_SIZE = (25, 25)
|
|
|
|
class Net(nn.Module):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.conv1 = nn.Conv2d(1, 12, 5) # 24x24 -> 21x21
|
|
self.maxpool = nn.MaxPool2d(2, 2) # 21x21 -> 10x10
|
|
self.conv2 = nn.Conv2d(12, 16, 3) # 10x10 -> 8x8
|
|
self.conv3 = nn.Conv2d(16, 32, 3) # 8x8 -> 6x6
|
|
self.fc1 = nn.Linear(32 * 6 * 6, 120)
|
|
self.fc2 = nn.Linear(120, 48)
|
|
self.fc3 = nn.Linear(48, 2)
|
|
|
|
def forward(self, x):
|
|
x = F.relu(self.conv1(x))
|
|
x = self.maxpool(x)
|
|
x = F.relu(self.conv2(x))
|
|
x = F.relu(self.conv3(x))
|
|
#print(x.shape)
|
|
|
|
x = torch.flatten(x, 1)
|
|
|
|
x = F.relu(self.fc1(x))
|
|
x = F.relu(self.fc2(x))
|
|
x = F.softmax(self.fc3(x))
|
|
return x
|
|
|
|
def classify_stone_empty(s: list[np.ndarray], model: nn.Module) -> bool:
|
|
s = torch.Tensor(s)
|
|
s = s[:,:,:,0]
|
|
s = s[:,torch.newaxis,:,:].float() / 255.0
|
|
s = s.to(0)
|
|
|
|
return torch.argmax(model(s), dim=1)
|
|
|
|
|
|
def load_model(f: str) -> torch.nn.Module:
|
|
torch.autograd.set_grad_enabled(False)
|
|
model = Net()
|
|
model.load_state_dict(torch.load(f, weights_only=True))
|
|
model.eval()
|
|
model.to(0)
|
|
return model |