<filename>src/models/model.py
from torch import nn
import torch.nn.functional as F
import torch
class MyAwesomeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(256, 120)  # 5*5 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Sequential(nn.Linear(84, 10),nn.LogSoftmax(dim=1))
        self.dropout = nn.Dropout(p=0.2)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = self.dropout(F.max_pool2d(F.relu(self.conv1(x)), (2, 2)))
        #print('x shape', x.shape)
        # If the size is a square, you can specify with a single number
        x = self.dropout(F.max_pool2d(F.relu(self.conv2(x)), 2))
        #print('x shape2', x.shape)
        x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension
        #print('x shape3', x.shape)
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.dropout(F.relu(self.fc2(x)))
        x = self.fc3(x)
        return x
<filename>src/models/model.py
from torch import nn
import torch.nn.functional as F
import torch
class MyAwesomeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(256, 120)  # 5*5 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Sequential(nn.Linear(84, 10),nn.LogSoftmax(dim=1))
        self.dropout = nn.Dropout(p=0.2)

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = self.dropout(F.max_pool2d(F.relu(self.conv1(x)), (2, 2)))
        #print('x shape', x.shape)
        # If the size is a square, you can specify with a single number
        x = self.dropout(F.max_pool2d(F.relu(self.conv2(x)), 2))
        #print('x shape2', x.shape)
        x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension
        #print('x shape3', x.shape)
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.dropout(F.relu(self.fc2(x)))
        x = self.fc3(x)
        return x
# an affine operation: y = Wx + b # 5*5 from image dimension # Max pooling over a (2, 2) window #print('x shape', x.shape) # If the size is a square, you can specify with a single number #print('x shape2', x.shape) # flatten all dimensions except the batch dimension #print('x shape3', x.shape)