107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
import torch
|
|
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import torch.optim as optim
|
|
|
|
import torchvision
|
|
import torchvision.transforms as transforms
|
|
|
|
|
|
class FFNet(nn.Module):
|
|
"""
|
|
Feedforward neural network, virtual base class
|
|
"""
|
|
def __init__(self):
|
|
super(FFNet, self).__init__()
|
|
self.dnn_model = self.build_model()
|
|
|
|
def build_model(self):
|
|
""" build specific model """
|
|
raise NotImplementedError
|
|
|
|
def forward(self, x):
|
|
""" feed data forward and return result """
|
|
raise NotImplementedError
|
|
|
|
def train_model(self, trainloader, testloader, loss_fn, optimizer, num_epochs):
|
|
"""Train a model."""
|
|
|
|
def train_epoch(model, dataloader, loss_fn, optimizer):
|
|
"""Train a single epoch"""
|
|
raise NotImplementedError
|
|
|
|
def test_epoch(model, dataloader, loss_fn):
|
|
"""Test a single epoch"""
|
|
raise NotImplementedError
|
|
|
|
raise NotImplementedError
|
|
|
|
def make_predict(self, images):
|
|
""" predict labels for images """
|
|
raise NotImplementedError
|
|
|
|
def evaluate_model(self, testloader, n):
|
|
""" evaluation: return accuracy """
|
|
raise NotImplementedError
|
|
|
|
def predict_one(self, x):
|
|
raise NotImplementedError
|
|
|
|
|
|
class MLP(FFNet):
|
|
def __init__(self):
|
|
super(MLP, self).__init__()
|
|
|
|
def build_model(self):
|
|
""" build specific model """
|
|
raise NotImplementedError
|
|
|
|
def forward(self, x):
|
|
""" feed data forward and return result """
|
|
raise NotImplementedError
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# load train and test set
|
|
trainset = torchvision.datasets.MNIST(
|
|
'../data', train=True, download=True, transform=transforms.ToTensor())
|
|
testset = torchvision.datasets.MNIST(
|
|
'../data', train=False, download=True, transform=transforms.ToTensor())
|
|
print(f'Train #: {len(trainset)}; Test #: {len(testset)}')
|
|
# data iterator
|
|
dataiter = iter(torch.utils.data.DataLoader(trainset, batch_size=8, shuffle=False))
|
|
images, labels = next(dataiter)
|
|
print(f'Labels: {labels}; Batch shape: {images.size()}')
|
|
|
|
# construct model
|
|
model = MLP()
|
|
print(model)
|
|
|
|
# data loader: easier iteration
|
|
BATCH_SIZE = 256
|
|
NUM_WORKERS = 4
|
|
trainloader = torch.utils.data.DataLoader(
|
|
trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS)
|
|
testloader = torch.utils.data.DataLoader(
|
|
testset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS)
|
|
|
|
# train the model
|
|
loss_fn = nn.CrossEntropyLoss() # cross entropy
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=0.1) # SGD
|
|
NUM_EPOCH = 5
|
|
model.train_model(trainloader, testloader, loss_fn, optimizer, NUM_EPOCH)
|
|
|
|
# check prediction accuracy
|
|
print(f'Labels : {labels}')
|
|
print(f'Prediction: {model.make_predict(images)}')
|
|
print(f'Accuracy: {model.evaluate_model(testloader, len(testset)):.2f}')
|
|
|
|
# application: predict hand-written digit
|
|
from PIL import Image
|
|
image = Image.open('number6c.png')
|
|
image = transforms.ToTensor()(image).unsqueeze(0)
|
|
print(f'loaded image shape: {image.size()}')
|
|
print(f'Predicted: "{model.predict_one(image)}", Actual: "{6}"')
|
|
|