diff --git a/src/cnn.py b/src/cnn.py new file mode 100644 index 0000000..733478a --- /dev/null +++ b/src/cnn.py @@ -0,0 +1,24 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class CNN(nn.Module): + def __init__(self): + super(CNN, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.fc2 = nn.Linear(128, 64) + self.fc3 = nn.Linear(64, 10) + + def forward(self, x): + x = F.relu(self.conv1(x)) + x = F.max_pool2d(x, 2, 2) + x = F.relu(self.conv2(x)) + x = F.max_pool2d(x, 2, 2) + x = x.view(-1, 64 * 7 * 7) + x = F.relu(self.fc1(x)) + x = F.relu(self.fc2(x)) + x = self.fc3(x) + return F.log_softmax(x, dim=1)