|
| 1 | +import torch |
| 2 | +from torch import nn |
| 3 | +from torch.utils.data import DataLoader |
| 4 | +from torchvision import datasets |
| 5 | +from torchvision.transforms import ToTensor |
| 6 | + |
| 7 | +# Download training data from open datasets. |
| 8 | +training_data = datasets.FashionMNIST( |
| 9 | + root="data", |
| 10 | + train=True, |
| 11 | + download=True, |
| 12 | + transform=ToTensor(), |
| 13 | +) |
| 14 | + |
| 15 | +# Download test data from open datasets. |
| 16 | +test_data = datasets.FashionMNIST( |
| 17 | + root="data", |
| 18 | + train=False, |
| 19 | + download=True, |
| 20 | + transform=ToTensor(), |
| 21 | +) |
| 22 | + |
| 23 | +batch_size = 64 |
| 24 | + |
| 25 | +# Create data loaders. |
| 26 | +train_dataloader = DataLoader(training_data, batch_size=batch_size) |
| 27 | +test_dataloader = DataLoader(test_data, batch_size=batch_size) |
| 28 | + |
| 29 | +for X, y in test_dataloader: |
| 30 | + print(f"Shape of X [N, C, H, W]: {X.shape}") |
| 31 | + print(f"Shape of y: {y.shape} {y.dtype}") |
| 32 | + break |
| 33 | + |
| 34 | +device = ( |
| 35 | + torch.accelerator.current_accelerator().type |
| 36 | + if torch.accelerator.is_available() |
| 37 | + else "cpu" |
| 38 | +) |
| 39 | +print(f"Using {device} device") |
| 40 | + |
| 41 | + |
| 42 | +# Define model |
| 43 | +class NeuralNetwork(nn.Module): |
| 44 | + def __init__(self): |
| 45 | + super().__init__() |
| 46 | + self.flatten = nn.Flatten() |
| 47 | + self.linear_relu_stack = nn.Sequential( |
| 48 | + nn.Linear(28 * 28, 512), |
| 49 | + nn.ReLU(), |
| 50 | + nn.Linear(512, 512), |
| 51 | + nn.ReLU(), |
| 52 | + nn.Linear(512, 10), |
| 53 | + ) |
| 54 | + |
| 55 | + def forward(self, x): |
| 56 | + x = self.flatten(x) |
| 57 | + logits = self.linear_relu_stack(x) |
| 58 | + return logits |
| 59 | + |
| 60 | + |
| 61 | +model = NeuralNetwork().to(device) |
| 62 | +print(model) |
| 63 | + |
| 64 | +loss_fn = nn.CrossEntropyLoss() |
| 65 | +optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) |
| 66 | + |
| 67 | + |
| 68 | +def train(dataloader, model, loss_fn, optimizer): |
| 69 | + size = len(dataloader.dataset) |
| 70 | + model.train() |
| 71 | + for batch, (X, y) in enumerate(dataloader): |
| 72 | + X, y = X.to(device), y.to(device) |
| 73 | + |
| 74 | + # Compute prediction error |
| 75 | + pred = model(X) |
| 76 | + loss = loss_fn(pred, y) |
| 77 | + |
| 78 | + # Backpropagation |
| 79 | + loss.backward() |
| 80 | + optimizer.step() |
| 81 | + optimizer.zero_grad() |
| 82 | + |
| 83 | + if batch % 100 == 0: |
| 84 | + loss, current = loss.item(), (batch + 1) * len(X) |
| 85 | + print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]") |
| 86 | + |
| 87 | + |
| 88 | +def test(dataloader, model, loss_fn): |
| 89 | + size = len(dataloader.dataset) |
| 90 | + num_batches = len(dataloader) |
| 91 | + model.eval() |
| 92 | + test_loss, correct = 0, 0 |
| 93 | + with torch.no_grad(): |
| 94 | + for X, y in dataloader: |
| 95 | + X, y = X.to(device), y.to(device) |
| 96 | + pred = model(X) |
| 97 | + test_loss += loss_fn(pred, y).item() |
| 98 | + correct += (pred.argmax(1) == y).type(torch.float).sum().item() |
| 99 | + test_loss /= num_batches |
| 100 | + correct /= size |
| 101 | + print( |
| 102 | + f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n" |
| 103 | + ) |
| 104 | + |
| 105 | + |
| 106 | +epochs = 5 |
| 107 | +for t in range(epochs): |
| 108 | + print(f"Epoch {t + 1}\n-------------------------------") |
| 109 | + train(train_dataloader, model, loss_fn, optimizer) |
| 110 | + test(test_dataloader, model, loss_fn) |
| 111 | +print("Done!") |
0 commit comments