Created
October 28, 2025 15:50
-
-
Save sebington/0ec87907ad0a079e6f59d941336d547c to your computer and use it in GitHub Desktop.
Code from the following YT video : https://www.youtube.com/watch?v=rdFjMb6dTJQ
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # ========================================================================= | |
| # A simple PyTorch example that learns y = x^2 | |
| # ========================================================================= | |
| #!/usr/bin/env -S uv run | |
| # /// script | |
| # requires-python = ">=3.9" | |
| # dependencies = [ | |
| # "torch", | |
| # "matplotlib", | |
| # "numpy", | |
| # ] | |
| # /// | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| import matplotlib.pyplot as plt | |
| # Step 1: Make a small synthetic dataset: y = x^2 + noise | |
| torch.manual_seed(0) | |
| N = 300 | |
| x = torch.linspace(-3, 3, N).unsqueeze(1) # shape [N, 1] | |
| y_true = x**2 | |
| y = y_true + 0.3 * torch.randn_like(y_true) # add a little noise | |
| # Step 2: Simple shuffle + train/test split | |
| perm = torch.randperm(N) | |
| train_idx = perm[:int(0.8*N)] | |
| test_idx = perm[int(0.8*N):] | |
| x_tr, y_tr = x[train_idx], y[train_idx] | |
| x_te, y_te = x[test_idx], y[test_idx] | |
| # Step 3: Define a simple neural net: 1 hidden layer with Tanh | |
| # (Input: 1 -> Hidden: 32 -> Output: 1) | |
| class SimpleNeuralNetwork(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.fc1 = nn.Linear(1, 32) | |
| self.act = nn.Tanh() | |
| self.fc2 = nn.Linear(32, 1) | |
| def forward(self, x): | |
| return self.fc2(self.act(self.fc1(x))) | |
| # Step 4: Instantiate SimpleNeuralNetwork object | |
| model = SimpleNeuralNetwork() | |
| # Step 5: Define Loss + Optimizer | |
| criterion = nn.MSELoss() | |
| optimizer = optim.Adam(model.parameters(), lr=1e-2) | |
| # Step 6: Train SimpleNeuralNetwork in training loop | |
| epochs = 500 | |
| for ep in range(1, epochs+1): | |
| model.train() | |
| y_hat = model(x_tr) # forward | |
| loss = criterion(y_hat, y_tr) # compute MSE | |
| optimizer.zero_grad() | |
| loss.backward() # backprop | |
| optimizer.step() # update weights | |
| if ep % 100 == 0 or ep == 1: | |
| print(f"Epoch {ep:3d} | train MSE: {loss.item():.4f}") | |
| # Step 7: Evaluate SimpleNeuralNetwork on test data | |
| model.eval() | |
| with torch.no_grad(): | |
| y_pred_te = model(x_te) | |
| test_mse = torch.mean((y_pred_te - y_te)**2).item() | |
| print(f"Test MSE: {test_mse:.4f}") | |
| # Step 8: Visualize: noisy points, true curve, and model prediction | |
| with torch.no_grad(): | |
| xs = torch.linspace(-3.5, 3.5, 400).unsqueeze(1) | |
| ys_pred = model(xs) | |
| plt.scatter(x_tr, y_tr, s=12, label="Train (noisy)") | |
| plt.scatter(x_te, y_te, s=16, label="Test (noisy)") | |
| plt.plot(xs, xs**2, linewidth=2, label="True y=x^2") | |
| plt.plot(xs, ys_pred, linewidth=2, linestyle="--", label="NN prediction") | |
| plt.legend() | |
| plt.xlabel("x"); plt.ylabel("y") | |
| plt.title("Simple PyTorch Regression: Learn y = x^2") | |
| plt.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment