PyTorch Quickstart Model

Beam Element

                # Define model
                class NeuralNetwork(nn.Module):
                    def __init__(self):
                        super().__init__()
                        self.flatten = nn.Flatten()
                        self.linear_relu_stack = nn.Sequential(
                            nn.Linear(28*28, 512),
                            nn.ReLU(),
                            nn.Linear(512, 512),
                            nn.ReLU(),
                            nn.Linear(512, 10)
                        )

                    def forward(self, x):
                        x = self.flatten(x)
                        logits = self.linear_relu_stack(x)
                        return logits


                def train(dataloader, model, loss_fn, optimizer):
                    size = len(dataloader.dataset)
                    model.train()
                    for batch, (X, y) in enumerate(dataloader):
                    X, y = X.to(device), y.to(device)
            
                    # Compute prediction error
                    pred = model(X)
                    loss = loss_fn(pred, y)
            
                    # Backpropagation
                    loss.backward()
                    optimizer.step()
                    optimizer.zero_grad()
            
                    if batch % 100 == 0:
                        loss, current = loss.item(), (batch + 1) * len(X)
                        print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


                classes = [
                "T-shirt/top",
                "Trouser",
                "Pullover",
                "Dress",
                "Coat",
                "Sandal",
                "Shirt",
                "Sneaker",
                "Bag",
                "Ankle boot",
                ]
                    
                model.eval()
                x, y = test_data[0][0], test_data[0][1]
                with torch.no_grad():
                    x = x.to(device)
                    pred = model(x)
                    predicted, actual = classes[pred[0].argmax(0)], classes[y]
                    print(f'Predicted: "{predicted}", Actual: "{actual}"')