Test torch

torch_mnist.py

import torch
from torchvision import transforms, datasets
from torch import nn, optim
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=3)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(500, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(-1, 500)
        x = F.relu(self.fc1(x))
        x = nn.Sigmoid()(self.fc2(x))
        return x

trainset = datasets.MNIST(root='.', train=True, download=True,
               transform=transforms.Compose([transforms.ToTensor()]))
testset  = datasets.MNIST(root='.', train=False, download=True,
               transform=transforms.Compose([transforms.ToTensor()]))
trainloader = torch.utils.data.DataLoader(trainset, batch_size=512, shuffle=False, num_workers=16)
testsloader = torch.utils.data.DataLoader(trainset, batch_size=512, shuffle=False, num_workers=16)

device = torch.device("cuda")
model = Net().to(device)
optimizer = optim.Adam(model.parameters())

for epoch in range(5):
    total = 0
    correct = 0
    for batch_idx, (data, target) in enumerate(trainloader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

        with torch.no_grad():
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{}]\tLoss: {:.6f}\tAcc: {:.6f}'.format(
                epoch, batch_idx * len(data), len(trainloader.dataset),
                loss.item(), 100 * correct / total))

    with torch.no_grad():
        total = 0
        correct = 0
        for data in testsloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        print(100 * correct / total)