CNN on Fashion-MNIST dataset with Pytorch

Fashion_MNST.py

import argparse
import torch 
import torch.nn as nn 
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.optim
import torch.utils.data

CUDA = False

#----------------------------- NET CLASS -------------------------------#

class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    #defining the layers
    self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=5) # 28 X 28 --> 24 X 24 --> 12 X 12
    self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=5) #12 X 12 --> 8 X 8 --> 4 X 4  
    self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=64) # 4 X 4 X 32 --> 64  
    self.fc2 = nn.Linear(in_features=64, out_features=10) # 64 --> 10

  #forward pass
  def forward(self, x):
    x = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=2, stride=2)
    x = F.max_pool2d(F.relu(self.conv2(x)), kernel_size=2, stride=2)
    x = x.view(-1, 16 * 4 * 4 )
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return x
  
def main(params):
  #get the train data download it if needed
  train = datasets.FashionMNIST(
          root = './data/FashionMNIST',
          train = True,
          download = True,
          transform = transforms.Compose([ transforms.ToTensor()]))
  
  #get the test data download it if needed
  test =  datasets.FashionMNIST(
          root = './data/FashionMNIST',
          train = False,
          download = True,
          transform = transforms.Compose([ transforms.ToTensor()]))
  #devide the train data into batchs and shaffle them 
  trainset = torch.utils.data.DataLoader(train, batch_size = params.batch_size,
                                           shuffle = True, num_workers = 2)
  
  #devide the train data into batchs and shaffle them 
  testset = torch.utils.data.DataLoader(test, batch_size = params.batch_size,
                                           shuffle = True, num_workers = 2)
  
  net = Net() #creating the net 
  criterion = nn.CrossEntropyLoss()
  optimizer = torch.optim.SGD(net.parameters(), params.lr) #creating the optimizer 
  
  if CUDA: #IF CUDA is enabled transforme the model into the GPU 
        net = net.cuda()
        criterion = criterion.cuda()
  
  #train the model 
  for i in range(params.epochs):
    print("=================\n=== EPOCH "+str(i+1)+" =====\n=================\n")
    for data in trainset:
      X, y = data
      if CUDA:
        X = X.cuda()
        y = y.cuda()
      net.zero_grad()
      output = net(X)
      loss = criterion(output, y)
      loss.backward()
      optimizer.step()
    print(loss)
  
  #test the model 
  correct = 0
  total = 0
  with torch.no_grad():
    for data in testset:
        X, y = data
        if CUDA:
          X = X.cuda()
          y = y.cuda()
        output = net(X)
        #print(output)
        for idx, i in enumerate(output):
            #print(torch.argmax(i), y[idx])
            if torch.argmax(i) == y[idx]:
                correct += 1
            total += 1
  print("Accuracy: ", round(correct/total, 3))
  
if __name__ == '__main__':

    # Paramètres en ligne de commande
    parser = argparse.ArgumentParser()
    #parser.add_argument('--path', default='/tmp/datasets/mnist', type=str, metavar='DIR', help='path to dataset')
    parser.add_argument('--epochs', default=5, type=int, metavar='N', help='number of total epochs to run')
    parser.add_argument('--batch-size', default=128, type=int, metavar='N', help='mini-batch size (default: 128)')
    parser.add_argument('--lr', default=0.1, type=float, metavar='LR', help='learning rate')
    parser.add_argument('--cuda', dest='cuda', action='store_true', help='activate GPU acceleration')

    args = parser.parse_args()
    if args.cuda:
        CUDA = True
        #cudnn.benchmark = True

    main(args)