Kuzushiji MNIST

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import sklearn.metrics as metrics
import numpy as np
from torchvision import datasets, transforms

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    conf_matrix = np.zeros((10,10)) # initialize confusion matrix
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            # sum up batch loss
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            # determine index with maximal log-probability
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            # update confusion matrix
            conf_matrix = conf_matrix + metrics.confusion_matrix(
                          target.cpu(),pred.cpu(),labels=[0,1,2,3,4,5,6,7,8,9])
        # print confusion matrix
        np.set_printoptions(precision=4, suppress=True)
        print(type(conf_matrix))
        print(conf_matrix)

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

class NetLin(nn.Module):
    # linear function followed by log_softmax
    def __init__(self):
        super(NetLin, self).__init__()
        self.fc = nn.Linear(28 * 28, 10)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # flatten the image into a vector
        x = self.fc(x)  # apply the linear layer
        x = self.log_softmax(x)  # apply log softmax
        return x

class NetFull(nn.Module):
    # two fully connected tanh layers followed by log softmax
    def __init__(self):
        super(NetFull, self).__init__()
        self.layer1 = nn.Linear(28 * 28, 100)
        self.layer2 = nn.Linear(100, 10)
        self.hid_act = nn.Tanh()
        self.hid_out = nn.LogSoftmax(dim=1)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = self.layer1(x)
        x = self.hid_act(x)
        x = self.layer2(x)
        x = self.hid_out(x)
        return x

class NetConv(nn.Module):
    # two convolutional layers and one fully connected layer,
    # all using relu, followed by log_softmax
    def __init__(self):
        super(NetConv, self).__init__()
        self.conv1 = nn.Conv2d(1, 128, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(128, 256, kernel_size=5, stride=2, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(256 * 6 * 6, 10)
        self.relu = nn.ReLU()
        self.act = nn.LogSoftmax(dim=1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool(x)

        #x = x.flatten(1)
        x = x.view(-1, 256 * 6 * 6)
        x = self.fc(x)
        x = self.act(x)

        return x

    def main(model, lr=0.01, mom=0.5, epochs=10):
      use_mps = torch.backends.mps.is_available()
      device = torch.device('mps' if use_mps else 'cpu')

      kwargs = {'num_workers': 1, 'pin_memory': True} if use_mps else {}

      # define a transform to normalize the data
      transform = transforms.Compose([transforms.ToTensor(),
				      transforms.Normalize((0.5,), (0.5,))])

      # fetch and load training data
      trainset = datasets.KMNIST(root='./data', train=True, download=True, transform=transform)
      train_loader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=False)

      # fetch and load test data
      testset = datasets.KMNIST(root='./data', train=False, download=True, transform=transform)
      test_loader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

      # choose network architecture
      if model == 'lin':
	  net = NetLin().to(device)
      elif model == 'full':
	  net = NetFull().to(device)
      else:
	  net = NetConv().to(device)

      if list(net.parameters()):
	  # use SGD optimizer
	  #optimizer = optim.SGD(net.parameters(), lr=lr, momentum=mom)

	  # use Adam optimizer
	  #optimizer = optim.Adam(net.parameters(),lr=lr,
	  #                             weight_decay=0.00001)
	  optimizer = optim.SGD(net.parameters(),lr=lr,momentum=0.9,
				      weight_decay=0.00001)

	  # training and testing loop
	  for epoch in range(1, epochs + 1):
	      train(net, device, train_loader, optimizer, epoch)
	      test(net, device, test_loader)

  main('lin')
  #main('full')
  #main('conv')