tanner/doormind
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add dropout and weight decay to prevent overfitting

Browse Source 
Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>
This commit is contained in:
Tanner Collin
2025-07-31 17:18:40 -06:00
parent c8f57818d1
commit 05ac4be541
2 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
model.py
View File

@@ -67,6 +67,7 @@ class GarageDoorCNN(nn.Module):
self.fc1 = nn.Linear(self.fc1_input_features, 512) self.fc1 = nn.Linear(self.fc1_input_features, 512)
self.relu4 = nn.ReLU() 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 self.fc2 = nn.Linear(512, 2) # 2 classes: open, closed
def forward(self, x): def forward(self, x):
@@ -75,5 +76,6 @@ class GarageDoorCNN(nn.Module):
x = self.pool3(self.relu3(self.conv3(x))) x = self.pool3(self.relu3(self.conv3(x)))
x = x.view(-1, self.fc1_input_features) # Flatten the tensor x = x.view(-1, self.fc1_input_features) # Flatten the tensor
x = self.relu4(self.fc1(x)) x = self.relu4(self.fc1(x))
x = self.dropout(x) # Apply dropout before the final layer
x = self.fc2(x) x = self.fc2(x)
return x return x

3
train.py
View File

@@ -67,6 +67,7 @@ def train_model():
NUM_EPOCHS = 10 NUM_EPOCHS = 10
BATCH_SIZE = 32 BATCH_SIZE = 32
LEARNING_RATE = 0.001 LEARNING_RATE = 0.001
WEIGHT_DECAY = 1e-5 # L2 regularization
# --- Data Preparation --- # --- Data Preparation ---
# Define separate transforms for training (with augmentation) and validation (without) # 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) model = GarageDoorCNN(resize_dim=RESIZE_DIM).to(device)
criterion = nn.CrossEntropyLoss() 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 --- # --- Training Loop ---
print("Starting training...") print("Starting training...")