Pytorch Neural Network Class Example
Full Script
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
# 1. Define a Custom Dataset Class
class RandomClassificationDataset(Dataset):
def __init__(self, num_samples, input_size, num_classes):
# Generate dummy data (X) and labels (y)
self.X = torch.randn(num_samples, input_size)
self.y = torch.randint(0, num_classes, (num_samples,))
def __len__(self):
# Must return the total number of samples
return len(self.X)
def __getitem__(self, idx):
# Must return a single sample (features, label) at the given index
return self.X[idx], self.y[idx]
# 2. Define the Neural Network Architecture
class SimpleClassifier(nn.Module):
def __init__(self, input_size, hidden_size, num_classes):
super(SimpleClassifier, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
# 3. Hyperparameters
INPUT_SIZE = 10
HIDDEN_SIZE = 32
NUM_CLASSES = 2
BATCH_SIZE = 16 # Number of samples processed before updating weights
LEARNING_RATE = 0.01
EPOCHS = 5
# 4. Prepare Dataset and DataLoader
# Instantiate the dataset with 200 total samples
train_dataset = RandomClassificationDataset(num_samples=200, input_size=INPUT_SIZE, num_classes=NUM_CLASSES)
# Wrap the dataset in a DataLoader
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
shuffle=True, # Shuffle data every epoch to reduce overfitting
drop_last=False # Keep the final batch even if it's smaller than BATCH_SIZE
)
# 5. Instantiate Model, Loss, and Optimizer
model = SimpleClassifier(INPUT_SIZE, HIDDEN_SIZE, NUM_CLASSES)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
# 6. Training Loop (with mini-batches)
print("Starting Training...")
for epoch in range(EPOCHS):
running_loss = 0.0
# Iterate over mini-batches provided by the DataLoader
for batch_idx, (batch_X, batch_y) in enumerate(train_loader):
# Forward pass
outputs = model(batch_X)
loss = criterion(outputs, batch_y)
# Backward pass and optimization
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
# Calculate average loss across all batches in this epoch
epoch_loss = running_loss / len(train_loader)
print(f"Epoch [{epoch+1}/{EPOCHS}], Average Loss: {epoch_loss:.4f}")
print("Training Complete!")