78 lines
2.1 KiB
Python
78 lines
2.1 KiB
Python
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
from torch.utils.data import DataLoader
|
||
|
|
from torchvision import datasets
|
||
|
|
from torchvision.transforms import v2
|
||
|
|
import datetime
|
||
|
|
|
||
|
|
time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
|
|
|
||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
|
print(device)
|
||
|
|
transform = v2.Compose([v2.ToImage() , v2.ToDtype(torch.float32, scale=True)])
|
||
|
|
# Download training data from open datasets.
|
||
|
|
training_data = datasets.FashionMNIST(
|
||
|
|
root="data",
|
||
|
|
train=True,
|
||
|
|
download=True,
|
||
|
|
transform=transform,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Download test data from open datasets.
|
||
|
|
test_data = datasets.FashionMNIST(
|
||
|
|
root="data",
|
||
|
|
train=False,
|
||
|
|
download=True,
|
||
|
|
transform=transform,
|
||
|
|
)
|
||
|
|
|
||
|
|
train_loader = DataLoader(training_data, batch_size=64, shuffle=True)
|
||
|
|
test_loader = DataLoader(test_data, batch_size=64, shuffle=False)
|
||
|
|
|
||
|
|
class NeuralNetwork(nn.Module):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.flatten = nn.Flatten()
|
||
|
|
self.network = nn.Sequential(
|
||
|
|
nn.Linear(784, 512),
|
||
|
|
nn.ReLU(),
|
||
|
|
nn.Linear(512, 512),
|
||
|
|
nn.ReLU(),
|
||
|
|
nn.Linear(512, 10),
|
||
|
|
)
|
||
|
|
def forward(self, x):
|
||
|
|
x = self.flatten(x)
|
||
|
|
return self.network(x)
|
||
|
|
|
||
|
|
model = NeuralNetwork().to(device)
|
||
|
|
print(model)
|
||
|
|
loss_fn = nn.CrossEntropyLoss()
|
||
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
|
||
|
|
|
||
|
|
epochs = 10
|
||
|
|
for epoch in range(epochs):
|
||
|
|
model.train()
|
||
|
|
for X, y in train_loader:
|
||
|
|
X = X.to(device)
|
||
|
|
y = y.to(device)
|
||
|
|
pred = model(X)
|
||
|
|
loss = loss_fn(pred, y)
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
optimizer.zero_grad()
|
||
|
|
print("Epoch:", epoch+1,"Loss:", loss.item())
|
||
|
|
|
||
|
|
model.eval()
|
||
|
|
correct = 0
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
for X, y in test_loader:
|
||
|
|
X = X.to(device)
|
||
|
|
y = y.to(device)
|
||
|
|
pred = model(X)
|
||
|
|
correct += (pred.argmax(1) == y).float().sum().item()
|
||
|
|
print("Accuracy:", correct/len(test_loader))
|
||
|
|
|
||
|
|
torch.save(model.state_dict(), f"fashion_mnist_epochs{epochs}_{time}.pth")
|
||
|
|
|
||
|
|
print("Model saved")
|