torch-randomdata-example.py

   1 import torch
   2 import torch.nn as nn
   3 import torch.optim as optim
   4 from torch.utils.data import Dataset, DataLoader
   5 import numpy as np
   6
   7
   8 # ----------------------------------- BUILD MODEL ---------------------------------------------
   9 # 1. Define the Model Class
  10 class SimpleDNN(nn.Module):
  11     def __init__(self, input_size, hidden_size, output_size):
  12         super(SimpleDNN, self).__init__()
  13         # Define the layers (Linear means fully connected)
  14         self.layer1 = nn.Linear(input_size, hidden_size)
  15         self.relu = nn.ReLU() # Activation function
  16         self.layer2 = nn.Linear(hidden_size, output_size)
  17
  18     def forward(self, x):
  19         # This is the path data takes through the network
  20         x = self.layer1(x)
  21         x = self.relu(x)
  22         x = self.layer2(x)
  23         return x
  24
  25 # Initialization: Assuming input is 784 (like MNIST flattened image)
  26 INPUT_SIZE = 784
  27 HIDDEN_SIZE = 128
  28 OUTPUT_SIZE = 10 # 10 classes
  29 model = SimpleDNN(INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE)
  30
  31
  32 # ----------------------------------- GPU INTEGRATION ---------------------------------------------
  33 # 1. Check for GPU availability
  34 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  35 print(f"Using device: {device}")
  36
  37 # 2. Move the entire model to the GPU
  38 model.to(device)
  39
  40 # Setup
  41 LEARNING_RATE = 0.001
  42 NUM_EPOCHS = 10
  43 criterion = nn.CrossEntropyLoss() # Loss function
  44 optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE) # Optimizer
  45
  46
  47 # ----------------------------------- LOAD DATA ---------------------------------------------
  48 # 1. DEFINE THE CUSTOM DATASET
  49 class CustomDataset(Dataset):
  50     def __init__(self, features, labels):
  51         # features should be the full dataset of inputs (e.g., all 784 pixel values)
  52         self.features = torch.tensor(features, dtype=torch.float32)
  53         # labels should be the full dataset of target labels (integers)
  54         self.labels = torch.tensor(labels, dtype=torch.long)
  55
  56     def __len__(self):
  57         # Returns the total number of samples
  58         return len(self.features)
  59
  60     def __getitem__(self, idx):
  61         # Returns a single sample and its label (formatted as a dictionary
  62         # to match your current usage: data['features'], data['labels'])
  63         return {
  64             'features': self.features[idx],
  65             'labels': self.labels[idx]
  66         }
  67
  68 # 2. LOAD THE DATA (REPLACE THIS WITH YOUR ACTUAL LOADING CODE)
  69 DUMMY_FEATURES = np.random.rand(100, INPUT_SIZE).astype(np.float32)
  70 DUMMY_LABELS = np.random.randint(0, OUTPUT_SIZE, 100).astype(np.int64)
  71
  72 # 3. INSTANTIATE AND WRAP THE LOADER
  73 # Create the dataset object
  74 train_dataset = CustomDataset(DUMMY_FEATURES, DUMMY_LABELS)
  75
  76 # Create the DataLoader object
  77 BATCH_SIZE = 64 # Choose a desired batch size
  78 train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
  79
  80
  81 # ----------------------------------- TRAINING ---------------------------------------------
  82 # --- The Training Loop ---
  83 for epoch in range(NUM_EPOCHS):
  84     for batch_idx, data in enumerate(train_loader):
  85
  86         # 1. MOVE DATA TO GPU
  87         inputs = data['features'].to(device)
  88         labels = data['labels'].to(device)
  89
  90         # 2. ZERO GRADIENTS (Crucial step!)
  91         # Must clear the gradients from the previous step
  92         optimizer.zero_grad()
  93
  94         # 3. FORWARD PASS
  95         outputs = model(inputs)
  96
  97         # 4. CALCULATE LOSS
  98         loss = criterion(outputs, labels)
  99
 100         # 5. BACKWARD PASS (Calculates gradients)
 101         # This is the step that utilizes CUDA for massive parallel computation.
 102         loss.backward()
 103
 104         # 6. OPTIMIZER STEP (Updates weights)
 105         optimizer.step()
 106
 107     print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Loss: {loss.item():.4f}")