ContLearn/run.py at main · krm9c/ContLearn · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
#!/usr/bin/env python
"""
Entry point for continual learning experiments.

Usage:
    python scripts/run.py <config_file>
    python scripts/run.py config/sine.json
    python scripts/run.py config/sine.json --runs 0
    python scripts/run.py config/sine.json --runs 2 --no-plots

Arguments:
    config_file: Path to JSON configuration file
    --runs: Run ID for this experiment (default: 0)
    --no-plots: Skip plot generation (default: generate plots)
    --figures-dir: Output directory for figures (default: figures)
"""
import argparse
import sys
import os
import warnings

# Add src to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
# Add run_files/scripts to path for plot_results import
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'run_files', 'scripts'))

import jax
from cl.config import Params, load_config
from cl.runners import train_model
from cl.core.recording import RecordingMixin

# from runs__.analysis.scripts.additional_python_scripts.scripts.plot_results import generate_plots

# Suppress JAX MaxPool gradient warning (expected behavior for second-order derivatives)
warnings.filterwarnings('ignore', message='.*reduce-window min/max.*')


def main():
    parser = argparse.ArgumentParser(
        description='Continual Learning Framework',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python run.py kkt_run/config/sine.json
    python run.py kkt_run/config/sine.json --runs 0
    python run.py kkt_run/config/sine.json --runs 2 --no-plots
    python run.py kkt_run/config/sine.json --runs 1 --figures-dir outputs/figures
        """
    )
    parser.add_argument('config', type=str, help='Path to JSON configuration file')
    parser.add_argument('--runs', type=int, default=0, help='Run ID for this experiment (default: 0)')
    parser.add_argument('--no-plots', action='store_true', help='Skip plot generation')
    parser.add_argument('--figures-dir', type=str, default='figures', help='Output directory for figures')
    parser.add_argument('--output-dir', type=str, default=None, help='Custom output directory for results')
    parser.add_argument('--model-suffix', type=str, default=None, help='Custom suffix for model/records files')

    args = parser.parse_args()

    # Load configuration
    if not os.path.exists(args.config):
        print(f"Error: Config file not found: {args.config}")
        sys.exit(1)

    config = load_config(args.config)

    print(f"JAX Backend: {jax.default_backend()}")
    print(f"JAX Devices: {jax.devices()}")
    print()

    print(f"Loaded config from: {args.config}")
    print(f"Problem type: {config.get('prob', 'unknown')}")
    print(f"Dataset: {config.get('data', 'unknown')}")
    print(f"Tasks: {config.get('n_task', 'unknown')}")
    print(f"AWB enabled: {config.get('awb_enabled', False)}")
    print()

    #Construct base path with AWB suffix if enabled and not already present
    if args.output_dir:
        # Use custom output directory (for Polaris runs)
        base_model_path = os.path.join(args.output_dir, args.model_suffix or 'model')
    else:
        # Default behavior
        base_model_path = config.get('model_path', 'outputs/model')
        if config.get('awb_enabled', False) and '_awb' not in base_model_path:
            base_model_path = f"{base_model_path}_awb"

    # Override config paths if custom output directory specified
    if args.output_dir:
        config['model_path'] = base_model_path
        print(f"Output directory: {args.output_dir}")
        print(f"Model path: {base_model_path}")
        print()

    run_id = args.runs

    print(f"\n{'#'*60}")
    print(f"# Run ID: {run_id}")
    print(f"{'#'*60}")

    # Generic unified runner works for all problem types
    record_dict = train_model(config, run_id=run_id)

    # # Generate plots for each run (unless --no-plots)
    # # Use AWB-aware figures directory
    # if not args.no_plots:
    #     figures_dir = args.figures_dir if args.figures_dir != 'figures' or args.output_dir else args.figures_dir
    #     # If using default figures dir and no custom output, add AWB suffix
    #     if figures_dir == 'figures' and not args.output_dir and config.get('awb_enabled', False):
    #         dataset_name = config.get('data', 'unknown')
    #         figures_dir = f'figures/{dataset_name}_awb'
    #     generate_plots(record_dict, output_dir=figures_dir, run_id=str(run_id))

    print("\nTraining complete!")


if __name__ == '__main__':
    main()