Summary

Add integration with Weights & Biases (WandB) for improved experiment tracking, visualization, and collaboration.

Description

TimeMixer would benefit from built-in WandB integration to help researchers track experiments more effectively. Currently, the codebase lacks a sophisticated experiment tracking solution, making it difficult to compare different model configurations and hyperparameters systematically.

Functionality to implement:

• Add WandB initialization in experimental setup
• Log model architecture, hyperparameters, and training configuration
• Track metrics during training and evaluation
• Log sample forecasts as visualizations
• Support experiment grouping for hyperparameter studies
• Add config flag to enable/disable WandB logging

Benefits:

• Improved experiment tracking and visualization
• Easier comparison between different model configurations
• Better collaboration among researchers
• Simplified hyperparameter tuning

Implementation approach

1. Add WandB as an optional dependency in requirements.txt
2. Create a wrapper class for logging in utils/loggers.py
3. Modify experiment classes to use the logger
4. Add command-line arguments to enable/disable logging

Related components

• exp/exp_long_term_forecasting.py
• exp/exp_short_term_forecasting.py
• Other experiment classes
• run.py (for adding command-line args)