53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader
|
|
from torchvision import datasets
|
|
from torchvision.transforms import Compose, Lambda, ToTensor
|
|
|
|
torch.manual_seed(10)
|
|
|
|
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
|
|
print(device)
|
|
|
|
trainset = datasets.FashionMNIST('data', train=True, download=True,
|
|
transform=Compose([ToTensor(), Lambda(lambda x: x.flatten())]))
|
|
testset = datasets.FashionMNIST('data', train=False, download=True,
|
|
transform=Compose([ToTensor(), Lambda(lambda x: x.flatten())]))
|
|
|
|
|
|
def train(model, dataset, n_iter=10, batch_size=256):
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
|
|
criterion = nn.NLLLoss()
|
|
dl = DataLoader(dataset, batch_size=batch_size)
|
|
model.train()
|
|
for epoch in range(n_iter):
|
|
for images, targets in dl:
|
|
optimizer.zero_grad()
|
|
out = model(images.to(device))
|
|
loss = criterion(out, targets.to(device))
|
|
loss.backward()
|
|
optimizer.step()
|
|
if epoch % 10 == 0:
|
|
print('epoch: %3d loss: %.4f' % (epoch, loss))
|
|
|
|
|
|
def accuracy(model, dataset):
|
|
model.eval()
|
|
correct = sum([(model(images.to(device)).argmax(dim=1) == targets.to(device)).sum()
|
|
for images, targets in DataLoader(dataset, batch_size=256)])
|
|
print(correct.float() / len(dataset))
|
|
|
|
|
|
neurons = 300
|
|
|
|
model = nn.Sequential(
|
|
nn.Linear(28 * 28, neurons),
|
|
nn.ReLU(),
|
|
nn.Linear(neurons, 10),
|
|
nn.LogSoftmax(dim=-1)
|
|
).to(device)
|
|
|
|
train(model, trainset)
|
|
|
|
accuracy(model, testset)
|