Skip to content

Conversation

@sridhs21
Copy link
Contributor

@sridhs21 sridhs21 commented Nov 11, 2025

Add Regularization and Data Augmentation

This PR adds regularization techniques and on-the-fly data augmentation to improve model training.

Changes

New Command-Line Arguments

  • --dropoutRate (default: 0.3) - Control dropout strength
  • --weightDecay (default: 5e-4) - Control L2 regularization
  • --seed (default: None) - Set random seed for reproducibility
  • --require-gpu - Exit if GPU not available

Dropout Improvements

  • UNet now accepts dropout_rate parameter
  • Added dropout layers after each decoder block (dec1_dropout, dec2_dropout, dec3_dropout, dec4_dropout)
  • Added dropout after bottleneck (bottleneck_dropout)

Weight Decay

  • Changed from hardcoded weight_decay=1e-5 to configurable weight_decay=args.weightDecay
  • Logged in optimizer initialization output

On-the-Fly Data Augmentation

  • Added augment parameter to XPointPatchDataset
  • Added seed parameter to XPointPatchDataset for reproducible augmentation
  • New _apply_augmentation() method with:
    • Random rotation (90°, 180°, 270°) - 75% probability
    • Random horizontal flip - 50% probability
    • Random vertical flip - 50% probability
    • Gaussian noise injection - 30% probability
    • Brightness/contrast adjustment - 30% probability
    • Random cutout/erasing - 20% probability
  • Training uses augment=True, validation uses augment=False
  • Removed static augmentation: XPointDataset(..., rotateAndReflect=False)

Reproducibility

  • New set_seed() function sets seeds for Python, NumPy, and PyTorch
  • Configures PyTorch deterministic mode when seed is provided
  • Seeds propagated to patch datasets

GPU Checking

  • --require-gpu flag exits with error if CUDA unavailable
  • Provides diagnostic messages for troubleshooting

@sridhs21 sridhs21 closed this Nov 11, 2025
@sridhs21 sridhs21 reopened this Nov 11, 2025
Copy link
Contributor

@cwsmith cwsmith left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is one small comment below. In addition to that, we discussed offline about changing the default checkpoint rate to every 100 epochs. Each checkpoint was about 100MB which could quickly fill up disk space for someone who doesn't expect it.

@cwsmith
Copy link
Contributor

cwsmith commented Nov 11, 2025

@sridhs21 would you also please change the default checkpoint interval to 100 epochs?

Copy link
Contributor

@cwsmith cwsmith left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you.

@cwsmith cwsmith merged commit 0e86a49 into main Nov 12, 2025
1 check passed
@cwsmith cwsmith deleted the fix-overfitting branch November 12, 2025 02:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants