feat: add dropout and weight decay to prevent overfitting
Co-authored-by: aider (gemini/gemini-2.5-pro-preview-05-06) <aider@aider.chat>
This commit is contained in:
3
train.py
3
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...")
|
||||
|
Reference in New Issue
Block a user