126 lines
3.3 KiB
Python
126 lines
3.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torchsummary import summary
|
|
|
|
class ChessPredictModelBaby(nn.Module):
|
|
def __init__(self):
|
|
super(ChessPredictModelBaby, self).__init__()
|
|
self.model = nn.Sequential(
|
|
nn.Conv2d(in_channels=8, out_channels=4, kernel_size=3, padding=1),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(kernel_size=2, stride=2),
|
|
|
|
nn.Conv2d(in_channels=4, out_channels=8, kernel_size=3, padding=1),
|
|
nn.ReLU(),
|
|
nn.AdaptiveAvgPool2d(output_size=1),
|
|
|
|
nn.Flatten(),
|
|
nn.Linear(8, 16),
|
|
nn.ReLU(),
|
|
nn.Linear(16, 1),
|
|
nn.Tanh()
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.model(x.permute(0, 3, 1, 2))
|
|
return x
|
|
|
|
class ChessPredictModelS(nn.Module):
|
|
def __init__(self):
|
|
super(ChessPredictModelS, self).__init__()
|
|
self.model = nn.Sequential(
|
|
nn.Conv2d(in_channels=8, out_channels=32, kernel_size=3, padding=1),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(kernel_size=2, stride=2),
|
|
|
|
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(kernel_size=2, stride=2),
|
|
|
|
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
|
|
nn.ReLU(),
|
|
nn.AdaptiveAvgPool2d(output_size=1),
|
|
|
|
nn.Flatten(),
|
|
nn.Linear(128, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 1),
|
|
nn.Tanh()
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.model(x.permute(0, 3, 1, 2))
|
|
return x
|
|
|
|
class BasicBlock(nn.Module):
|
|
def __init__(self, in_channels, out_channels, stride=1):
|
|
super(BasicBlock, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
|
|
self.bn1 = nn.BatchNorm2d(out_channels)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
|
self.bn2 = nn.BatchNorm2d(out_channels)
|
|
|
|
self.downsample = None
|
|
if stride != 1 or in_channels != out_channels:
|
|
self.downsample = nn.Sequential(
|
|
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride),
|
|
nn.BatchNorm2d(out_channels)
|
|
)
|
|
|
|
def forward(self, x):
|
|
identity = x
|
|
|
|
out = self.conv1(x)
|
|
out = self.bn1(out)
|
|
out = self.relu(out)
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out)
|
|
|
|
if self.downsample is not None:
|
|
identity = self.downsample(x)
|
|
|
|
out += identity
|
|
out = self.relu(out)
|
|
|
|
return out
|
|
|
|
class ResNet(nn.Module):
|
|
def __init__(self, block, layers, num_classes=1):
|
|
super(ResNet, self).__init__()
|
|
self.in_channels = 64
|
|
|
|
self.model = nn.Sequential(
|
|
nn.Conv2d(8, 64, kernel_size=3, stride=1, padding=1),
|
|
nn.BatchNorm2d(64),
|
|
nn.ReLU(inplace=True),
|
|
self._make_layer(block, 64, layers[0], stride=1),
|
|
self._make_layer(block, 128, layers[1], stride=2),
|
|
self._make_layer(block, 256, layers[2], stride=2),
|
|
self._make_layer(block, 512, layers[3], stride=2),
|
|
nn.AdaptiveAvgPool2d((1, 1)),
|
|
nn.Flatten(),
|
|
nn.Linear(512, num_classes),
|
|
nn.Tanh()
|
|
)
|
|
|
|
def _make_layer(self, block, out_channels, blocks, stride=1):
|
|
layers = []
|
|
layers.append(block(self.in_channels, out_channels, stride))
|
|
self.in_channels = out_channels
|
|
for _ in range(1, blocks):
|
|
layers.append(block(out_channels, out_channels))
|
|
return nn.Sequential(*layers)
|
|
|
|
def forward(self, x):
|
|
return self.model(x.permute(0, 3, 1, 2))
|
|
|
|
def resnet18():
|
|
return ResNet(BasicBlock, [2, 2, 2, 2])
|
|
|
|
if __name__ == '__main__':
|
|
model = ChessPredictModelS()
|
|
summary(model, (8, 8, 8))
|