ICANSee/ensemble_train.py at master · Kirscher/ICANSee · 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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
#!/usr/bin/env python3
"""
Train multiple models with different seeds for ensemble prediction.
"""

import subprocess
import argparse


def train_single_model(seed, args):
    """Train a single model with given seed."""
    cmd = [
        "python3",
        "train.py",
        "--data_path",
        args.data_path,
        "--model",
        args.model,
        "--seed",
        str(seed),
        "--epochs",
        str(args.epochs),
        "--batch_size",
        str(args.batch_size),
        "--lr",
        str(args.lr),
        "--lr_backbone",
        str(args.lr_backbone),
        "--lr_head",
        str(args.lr_head),
        "--freeze_backbone_epochs",
        str(args.freeze_backbone_epochs),
        "--loss",
        args.loss,
        "--focal_alpha",
        str(args.focal_alpha),
        "--focal_gamma",
        str(args.focal_gamma),
        "--sampling",
        args.sampling,
        "--save_metric",
        args.save_metric,
        "--input_size",
        str(args.input_size),
    ]

    if args.pretrained_path:
        cmd.extend(["--pretrained_path", args.pretrained_path])
    elif args.use_gastronet:
        cmd.append("--use_gastronet")

    print(f"Training model with seed {seed}")
    print("Command:", " ".join(cmd))

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode == 0:
        print(f"Model with seed {seed} trained successfully")
        return True
    else:
        print(f"Model with seed {seed} failed:")
        print(result.stderr)
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Train ensemble of models with different seeds"
    )
    parser.add_argument(
        "--data_path", type=str, required=True, help="Path to dataset"
    )
    parser.add_argument(
        "--model", type=str, default="resnet50", help="Model architecture"
    )
    parser.add_argument(
        "--seeds",
        type=int,
        nargs="+",
        default=[42, 123, 456, 789, 999],
        help="Random seeds for ensemble",
    )
    parser.add_argument(
        "--epochs", type=int, default=30, help="Number of epochs"
    )
    parser.add_argument(
        "--batch_size", type=int, default=32, help="Batch size"
    )
    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
    parser.add_argument(
        "--lr_backbone",
        type=float,
        default=1e-5,
        help="Backbone learning rate",
    )
    parser.add_argument(
        "--lr_head", type=float, default=3e-4, help="Head learning rate"
    )
    parser.add_argument(
        "--freeze_backbone_epochs",
        type=int,
        default=5,
        help="Freeze backbone epochs",
    )
    parser.add_argument(
        "--loss", type=str, default="focal_loss", help="Loss function"
    )
    parser.add_argument(
        "--focal_alpha", type=float, default=0.9, help="Focal loss alpha"
    )
    parser.add_argument(
        "--focal_gamma", type=float, default=1.5, help="Focal loss gamma"
    )
    parser.add_argument(
        "--sampling", type=str, default="oversample", help="Sampling strategy"
    )
    parser.add_argument(
        "--save_metric", type=str, default="PPV@90% Recall", help="Save metric"
    )
    parser.add_argument(
        "--input_size", type=int, default=224, help="Input image size"
    )
    parser.add_argument(
        "--pretrained_path", type=str, help="Path to pretrained weights"
    )
    parser.add_argument(
        "--use_gastronet", action="store_true", help="Use Gastronet weights"
    )

    args = parser.parse_args()

    print(f"Training ensemble of {len(args.seeds)} models")
    print(f"Seeds: {args.seeds}")

    successful_models = []

    for seed in args.seeds:
        success = train_single_model(seed, args)
        if success:
            successful_models.append(seed)

    print("\nEnsemble training complete!")
    num_success = len(successful_models)
    num_total = len(args.seeds)
    print(f"Successfully trained {num_success}/{num_total} models")
    print(f"Successful seeds: {successful_models}")

    if successful_models:
        print("\nTo evaluate ensemble, run:")
        print(
            "python3 ensemble_evaluate.py --data_path",
            args.data_path,
            "--seeds",
            *successful_models,
        )


if __name__ == "__main__":
    main()