From 05ac4be54189a1efe68a0f34e4664c94dea81448 Mon Sep 17 00:00:00 2001 From: Tanner Collin Date: Thu, 31 Jul 2025 17:18:40 -0600 Subject: [PATCH] feat: add dropout and weight decay to prevent overfitting Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) --- model.py | 2 ++ train.py | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/model.py b/model.py index 86d563f..e35d5a7 100644 --- a/model.py +++ b/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 diff --git a/train.py b/train.py index 2af723d..4905855 100644 --- a/train.py +++ b/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...")