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!")