90 lines
2.3 KiB
Python
90 lines
2.3 KiB
Python
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader
|
|
from torchvision import datasets
|
|
from torchvision.transforms import ToTensor, Lambda, Compose
|
|
import matplotlib.pyplot as plt
|
|
|
|
# Download training data from open datasets.
|
|
training_data = datasets.FashionMNIST(
|
|
root="data",
|
|
train=True,
|
|
download=True,
|
|
transform=ToTensor(),
|
|
)
|
|
|
|
# Download test data from open datasets.
|
|
test_data = datasets.FashionMNIST(
|
|
root="data",
|
|
train=False,
|
|
download=True,
|
|
transform=ToTensor(),
|
|
)
|
|
|
|
batch_size = 64
|
|
|
|
# Create data loaders.
|
|
train_loader = DataLoader(training_data, batch_size=batch_size)
|
|
test_loader = DataLoader(test_data, batch_size=batch_size)
|
|
|
|
|
|
def output_label(label):
|
|
output_mapping = {
|
|
0: "T-shirt/Top",
|
|
1: "Trouser",
|
|
2: "Pullover",
|
|
3: "Dress",
|
|
4: "Coat",
|
|
5: "Sandal",
|
|
6: "Shirt",
|
|
7: "Sneaker",
|
|
8: "Bag",
|
|
9: "Ankle Boot"
|
|
}
|
|
input = (label.item() if type(label) == torch.Tensor else label)
|
|
return output_mapping[input]
|
|
|
|
class FashionCNN(nn.Module):
|
|
|
|
def __init__(self):
|
|
super(FashionCNN, self).__init__()
|
|
|
|
self.layer1 = nn.Sequential(
|
|
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1),
|
|
nn.BatchNorm2d(32),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(kernel_size=2, stride=2)
|
|
)
|
|
|
|
self.layer2 = nn.Sequential(
|
|
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
|
|
nn.BatchNorm2d(64),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(2)
|
|
)
|
|
|
|
self.fc1 = nn.Linear(in_features=64*6*6, out_features=600)
|
|
self.drop = nn.Dropout2d(0.25)
|
|
self.fc2 = nn.Linear(in_features=600, out_features=120)
|
|
self.fc3 = nn.Linear(in_features=120, out_features=10)
|
|
|
|
def forward(self, x):
|
|
out = self.layer1(x)
|
|
out = self.layer2(out)
|
|
out = out.view(out.size(0), -1)
|
|
out = self.fc1(out)
|
|
out = self.drop(out)
|
|
out = self.fc2(out)
|
|
out = self.fc3(out)
|
|
|
|
return out
|
|
|
|
model = FashionCNN()
|
|
|
|
error = nn.CrossEntropyLoss()
|
|
|
|
learning_rate = 0.001
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
|
|
|
|