import torch
from torch import nn, optim
from torch.nn import init
import torchvision
from torchvision import transforms
import numpy as np
import matplotlib.pyplot as plt
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
num_inputs = 28 * 28
num_outputs = 10
batch_size = 256
num_workers = 0
num_epochs = 10
lr = 0.5
class MLP(nn.Module):
def __init__(self, num_inputs, num_outputs):
super(MLP, self).__init__()
self.mlp = nn.Sequential(
nn.Linear(num_inputs, 256),
nn.ReLU(),
nn.Linear(256, 100),
nn.ReLU(),
nn.Linear(100, num_outputs)
)
def forward(self, x):
x = x.view(x.shape[0], -1)
x = self.mlp(x)
return x
def evaluate_accuracy(data_iter, net):
acc_sum, n = 0.0, 0
for X, y in data_iter:
X = X.to(device)
y = y.to(device)
acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
n += y.shape[0]
return acc_sum / n
def main():
root = "./Datasets/FashionMNIST"
mnist_train = torchvision.datasets.FashionMNIST(root=root, train=True, download=True,
transform=transforms.ToTensor())
mnist_test = torchvision.datasets.FashionMNIST(root=root, train=False, download=True,
transform=transforms.ToTensor())
train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)
net = MLP(num_inputs, num_outputs)
net = net.to(device)
print("training on ", device)
for params in net.parameters():
init.normal_(params, mean=0, std=0.01)
loss = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=lr)
loss_list, train_acc_list, test_acc_list = [], [], []
for epoch in range(num_epochs):
train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
for X, y in train_iter:
X = X.to(device)
y = y.to(device)
y_hat = net(X)
l = loss(y_hat, y).sum()
optimizer.zero_grad()
l.backward()
optimizer.step()
train_l_sum += l.item()
train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
n += y.shape[0]
test_acc = evaluate_accuracy(test_iter, net)
loss_list.append(train_l_sum / n)
train_acc_list.append(train_acc_sum / n)
test_acc_list.append(test_acc)
print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f'
% (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))
idx = range(1, num_epochs + 1)
plt.figure(figsize=(12, 4))
plt.subplot(121)
plt.plot(idx, loss_list, 'o-')
plt.grid(), plt.xlabel("epoch"), plt.ylabel("loss"), plt.xticks(range(min(idx), max(idx) + 1, 1))
plt.subplot(122)
plt.plot(idx, train_acc_list, 'ro-', label="train accuracy")
plt.plot(idx, test_acc_list, 'bo-', label="test accuracy")
plt.xlabel("epoch"), plt.ylabel("accuracy"), plt.legend(loc="best")
plt.xticks(range(min(idx), max(idx) + 1, 1))
plt.grid(), plt.ylim([0, 1.1])
plt.show()
if __name__ == '__main__':
main()
