feat: add dropout and weight decay to prevent overfitting

Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>

model.py

@@ -67,6 +67,7 @@ class GarageDoorCNN(nn.Module):
         
         self.fc1 = nn.Linear(self.fc1_input_features, 512)
         self.relu4 = nn.ReLU()
+        self.dropout = nn.Dropout(0.5) # Add dropout with 50% probability
         self.fc2 = nn.Linear(512, 2) # 2 classes: open, closed
 
     def forward(self, x):
@@ -75,5 +76,6 @@ class GarageDoorCNN(nn.Module):
         x = self.pool3(self.relu3(self.conv3(x)))
         x = x.view(-1, self.fc1_input_features) # Flatten the tensor
         x = self.relu4(self.fc1(x))
+        x = self.dropout(x) # Apply dropout before the final layer
         x = self.fc2(x)
         return x

train.py

@@ -67,6 +67,7 @@ def train_model():
     NUM_EPOCHS = 10
     BATCH_SIZE = 32
     LEARNING_RATE = 0.001
+    WEIGHT_DECAY = 1e-5 # L2 regularization
 
     # --- Data Preparation ---
     # Define separate transforms for training (with augmentation) and validation (without)
@@ -123,7 +124,7 @@ def train_model():
 
     model = GarageDoorCNN(resize_dim=RESIZE_DIM).to(device)
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
 
     # --- Training Loop ---
     print("Starting training...")