Skip to content

Commit

Permalink
feat: Added CNN class and training function in cnn
Browse files Browse the repository at this point in the history
  • Loading branch information
sweep-nightly[bot] authored Nov 26, 2023
1 parent 7284908 commit d5b3bff
Showing 1 changed file with 54 additions and 0 deletions.
54 changes: 54 additions & 0 deletions src/cnn.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=5)
self.conv2 = nn.Conv2d(32, 64, kernel_size=5)
self.dropout1 = nn.Dropout2d(0.25)
self.dropout2 = nn.Dropout2d(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)

def forward(self, x):
x = self.conv1(x)
x = nn.functional.relu(x)
x = self.conv2(x)
x = nn.functional.relu(x)
x = nn.functional.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
output = nn.functional.log_softmax(x, dim=1)
return output

def train(network, train_loader, optimizer):
network.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = network(data)
loss = nn.functional.nll_loss(output, target)
loss.backward()
optimizer.step()

if __name__ == "__main__":
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])

trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)

network = CNN()
optimizer = optim.SGD(network.parameters(), lr=0.01)

train(network, trainloader, optimizer)

0 comments on commit d5b3bff

Please sign in to comment.