1604671d36
Five modules covering nanoGPT, Ollama, RAG, semantic search, and neural networks. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
# nn_torch.py
|
|
#
|
|
# The same neural network as nn_numpy.py, but using PyTorch.
|
|
# Compare this to the numpy version to see what the framework handles for you:
|
|
# - Automatic differentiation (no manual backprop)
|
|
# - Built-in optimizers (Adam instead of hand-coded gradient descent)
|
|
# - GPU support (if available)
|
|
#
|
|
# CHEG 667-013
|
|
# E. M. Furst
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
# ── Load and prepare data ──────────────────────────────────────
|
|
|
|
data = np.loadtxt("data/n2_cp.csv", delimiter=",", skiprows=1)
|
|
T_raw = data[:, 0]
|
|
Cp_raw = data[:, 1]
|
|
|
|
# Normalize to [0, 1]
|
|
T_min, T_max = T_raw.min(), T_raw.max()
|
|
Cp_min, Cp_max = Cp_raw.min(), Cp_raw.max()
|
|
|
|
X = torch.tensor((T_raw - T_min) / (T_max - T_min), dtype=torch.float32).reshape(-1, 1)
|
|
Y = torch.tensor((Cp_raw - Cp_min) / (Cp_max - Cp_min), dtype=torch.float32).reshape(-1, 1)
|
|
|
|
# ── Define the network ─────────────────────────────────────────
|
|
#
|
|
# nn.Sequential stacks layers in order. Compare this to nanoGPT's
|
|
# model.py, which uses the same PyTorch building blocks (nn.Linear,
|
|
# activation functions) but with many more layers.
|
|
|
|
H = 10 # hidden neurons
|
|
|
|
model = nn.Sequential(
|
|
nn.Linear(1, H), # input -> hidden (W1, b1)
|
|
nn.Tanh(), # activation
|
|
nn.Linear(H, 1), # hidden -> output (W2, b2)
|
|
)
|
|
|
|
print(f"Model:\n{model}")
|
|
print(f"Total parameters: {sum(p.numel() for p in model.parameters())}\n")
|
|
|
|
# ── Training ───────────────────────────────────────────────────
|
|
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
|
|
loss_fn = nn.MSELoss()
|
|
|
|
epochs = 5000
|
|
log_interval = 500
|
|
losses = []
|
|
|
|
for epoch in range(epochs):
|
|
# Forward pass -- PyTorch tracks operations for automatic differentiation
|
|
Y_pred = model(X)
|
|
loss = loss_fn(Y_pred, Y)
|
|
losses.append(loss.item())
|
|
|
|
# Backward pass -- PyTorch computes all gradients automatically
|
|
optimizer.zero_grad() # reset gradients from previous step
|
|
loss.backward() # compute gradients via automatic differentiation
|
|
optimizer.step() # update weights (Adam optimizer)
|
|
|
|
if epoch % log_interval == 0 or epoch == epochs - 1:
|
|
print(f"Epoch {epoch:5d} Loss: {loss.item():.6f}")
|
|
|
|
# ── Results ────────────────────────────────────────────────────
|
|
|
|
# Predict on a fine grid
|
|
T_fine = torch.linspace(0, 1, 200).reshape(-1, 1)
|
|
with torch.no_grad(): # no gradient tracking needed for inference
|
|
Cp_pred_norm = model(T_fine)
|
|
|
|
# Convert back to physical units
|
|
T_fine_K = T_fine.numpy() * (T_max - T_min) + T_min
|
|
Cp_pred = Cp_pred_norm.numpy() * (Cp_max - Cp_min) + Cp_min
|
|
|
|
# ── Plot ───────────────────────────────────────────────────────
|
|
|
|
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
|
|
|
|
ax1.plot(T_raw, Cp_raw, 'ko', markersize=6, label='NIST data')
|
|
ax1.plot(T_fine_K, Cp_pred, 'r-', linewidth=2, label=f'NN ({H} neurons)')
|
|
ax1.set_xlabel('Temperature (K)')
|
|
ax1.set_ylabel('$C_p$ (kJ/kg/K)')
|
|
ax1.set_title('$C_p(T)$ for N$_2$ at 1 bar — PyTorch')
|
|
ax1.legend()
|
|
|
|
ax2.semilogy(losses)
|
|
ax2.set_xlabel('Epoch')
|
|
ax2.set_ylabel('Mean Squared Error')
|
|
ax2.set_title('Training Loss')
|
|
|
|
plt.tight_layout()
|
|
plt.savefig('nn_fit_torch.png', dpi=150)
|
|
plt.show()
|