BlueDisc/02_inference.py at main · SeisBlue/BlueDisc · 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
import argparse
import os

import mlflow
import seisbench.data as sbd
import seisbench.generate as sbg
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from module.gan_model import GANModel
from module.logger import MLFlowLogger
from module.pipeline import AugmentationsBuilder
from module.random_seed import RandomSeedManager
from module.device_manager import DeviceManager

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--run-id", type=str, required=True, help="MLflow run ID for the trained model")
    parser.add_argument("--dataset", type=str, required=True, help="Dataset name available in seisbench dataset class name (e.g., ETHZ, InstanceCount)")
    parser.add_argument("--step", type=int, help="Checkpoint step to load")
    parser.add_argument("--epoch", type=int, help="Checkpoint epoch to load")
    parser.add_argument("--device", type=str, default="auto", help="Device to use for inference (e.g., 'cpu', 'cuda', 'auto')")
    parser.add_argument(
        "--data-split",
        type=str,
        default="test",
        choices=["track", "train", "dev", "test"],
    )

    args = parser.parse_args()

    seed_value = 42
    seed_manager = RandomSeedManager(seed_value)
    seed_manager.set_seed()

    # Initialize device manager
    device_manager = DeviceManager(args.device)
    print(f"Using device: {device_manager.device}")

    # MLflow settings - using local setup
    mlflow_host = "127.0.0.1"
    mlflow_port = 5000

    run_id = args.run_id
    client = mlflow.MlflowClient(f"http://{mlflow_host}:{mlflow_port}")
    experiment_id = client.get_run(run_id).info.experiment_id
    experiment_name = client.get_experiment(experiment_id).name

    # Initialize logger
    logger = MLFlowLogger(
        run_id=run_id,
        mlflow_host=mlflow_host,
        mlflow_port=mlflow_port
    )

    base_path = logger.base_path
    checkpoint_dir = os.path.join(base_path, "checkpoint")

    gan_model = GANModel()
    gan_model.load_checkpoint_by_dir(checkpoint_dir, step=args.step, epoch=args.epoch)

    # Move models to device
    g_model = device_manager.move_to_device(gan_model.g_model)

    if gan_model.d_model:
        gan_model.d_model = device_manager.move_to_device(gan_model.d_model)

    # Load dataset
    data_class = getattr(sbd, args.dataset)
    data = data_class(sampling_rate=100)
    train, dev, test = data.train_dev_test()

    aug_builder = AugmentationsBuilder(dataset=data)
    augmentations = aug_builder.build()

    test_generator = sbg.GenericGenerator(dev)
    test_generator.add_augmentations(augmentations)

    g = torch.Generator()
    g.manual_seed(seed_value)

    num_workers = os.cpu_count() or 1

    test_loader = DataLoader(
        test_generator,
        batch_size=1000,
        shuffle=False,
        num_workers=num_workers,
        worker_init_fn=seed_manager.worker_init_fn,
        pin_memory=True,
        drop_last=True,
        generator=g,
    )


    size = len(test_generator)
    batch_size = test_loader.batch_size
    gan_model.eval()
    with torch.no_grad():
        with tqdm(total=size, desc="Inferencing", ncols=80) as pbar:
            for batch_id, batch in enumerate(test_loader):
                if batch_id == size:
                    break
                trace_name_list = batch.pop("trace_name", "")
                logger.log_text(
                    trace_name_list,
                    data_split=args.data_split,
                    data_type="trace_name",
                    step=batch_id,
                )

                batch = {
                    k: v.to(torch.float32).to(g_model.device) for k, v in batch.items()
                }
                y_sample = g_model.reorder_label_phase(batch).to(g_model.device)

                g_pred = g_model(batch["X"], logits=g_model.logits)
                g_pred = torch.sigmoid(g_pred).to(g_model.device)

                g_pred = g_pred.detach().float().cpu().numpy()
                y_sample = y_sample.detach().float().cpu().numpy()
                batch["X"] = batch["X"].detach().float().cpu().numpy()

                logger.log_hdf5(
                    batch["X"],
                    data_split=args.data_split,
                    data_type="waveform",
                    step=batch_id,
                )
                logger.log_hdf5(
                    y_sample,
                    data_split=args.data_split,
                    data_type="label",
                    step=batch_id,
                )
                logger.log_hdf5(
                    g_pred,
                    data_split=args.data_split,
                    data_type="prediction",
                    step=batch_id,
                )

                pbar.set_postfix()
                pbar.update(batch_size)

    print(f"Inference complete for {args.data_split} split")