model.py

import torch.nn as nn
import torch.nn.functional as F

# conv1 out size = (20 - 3 + 0)/1 + 1 = 18
# pool1 out size = (18 - 2 + 0)/2 + 1 = 9
# conv2 out size = (9 - 2 + 0)/1 + 1 = 8
# pool2 out size = (8 - 2 + 0)/2 + 1 = 4
POOL2_OUTSIZE = 4


class Model(nn.Module):
    def __init__(self):
        """
        Init model to decide if the image is a cross, an equal or nothing.
        """
        super(Model, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=2)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.linear1 = nn.Linear(int(POOL2_OUTSIZE**2) * 16, 120)
        self.linear2 = nn.Linear(120, 84)
        self.linear3 = nn.Linear(84, 3)

    def forward(self, x):
        """
        x: matriz 20x20
        """
        # Max pooling over a (2, 2) window
        x = self.pool(F.relu(self.conv1(x)))
        # If the size is a square you can only specify a single number
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, int((POOL2_OUTSIZE**2) * 16))
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = F.softmax(self.linear3(x), dim=1)
        return x