PreSimNet/evaluate_position_rmse.py at main · HaoLyu666/PreSimNet · 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
from __future__ import annotations

import argparse
import csv
import os

import numpy as np
import torch as t
from tqdm import tqdm

from config import args as base_args
import model.presimnet as model


CHECKPOINTS = {
    "presimnet_top2": (
        "checkpoints/"
        "presimnet_top2/"
        "epoch21_e.tar"
    ),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate direct position prediction RMSE.")
    parser.add_argument("--data", default="../data/test_data.npy")
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--out", default="outputs/position_rmse_compare.csv")
    parser.add_argument("--device", default="cuda:0" if t.cuda.is_available() else "cpu")
    parser.add_argument("--drop-last", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--max-batches", type=int, default=None)
    return parser.parse_args()


def make_batch_iterator(data_path, batch_size, in_length, drop_last, max_batches):
    data = np.load(data_path, mmap_mode="r")
    total = len(data)
    usable = (total // batch_size) * batch_size if drop_last else total
    batch_count = (usable + batch_size - 1) // batch_size
    if max_batches is not None:
        batch_count = min(batch_count, max_batches)
        usable = min(usable, batch_count * batch_size)

    for start in range(0, usable, batch_size):
        end = min(start + batch_size, usable)
        window = data[start:end]
        hist_acc_diff = window[:, :, 5] - window[:, :, 7]
        hist = np.concatenate(
            [window[:, :in_length, 4:], hist_acc_diff[:, :in_length, None]],
            axis=2,
        ).astype(np.float32, copy=False)
        fut_pos = window[:, in_length:, 11:12].astype(np.float32, copy=False)
        yield hist, fut_pos


def eval_checkpoint(label, checkpoint_path, run_args, device, cli):
    encoder = model.Encoder(run_args).to(device)
    checkpoint = t.load(checkpoint_path, map_location=device)
    encoder.load_state_dict(checkpoint["model_state_dict"])
    encoder.eval()

    loss_vals = t.zeros(run_args["out_length"], 1, device=device)
    counts = t.zeros(run_args["out_length"], 1, device=device)
    samples = 0

    with t.no_grad():
        for hist_np, fut_pos_np in tqdm(
            make_batch_iterator(
                cli.data,
                cli.batch_size,
                run_args["in_length"],
                cli.drop_last,
                cli.max_batches,
            ),
            desc=f"Position RMSE {label}",
        ):
            hist = t.from_numpy(hist_np).to(device)
            fut_pos = t.from_numpy(fut_pos_np).to(device)
            pred = encoder(hist)["position_pred"]
            mask = t.ones_like(fut_pos)
            err = (pred - fut_pos).pow(2) * mask
            loss_vals += t.sqrt(t.sum(err, dim=0))
            counts += t.sqrt(t.sum(mask, dim=0))
            samples += hist.shape[0]

    rmse = (loss_vals / (counts + 1e-12)).detach().cpu().numpy().reshape(-1)
    horizons = [4, 9, 14, 19]
    row = {"model": label, "checkpoint": checkpoint_path, "n_samples": samples}
    for idx in horizons:
        row[f"pos_rmse_{(idx + 1) * 0.1:.1f}s"] = float(rmse[idx])
    row["pos_rmse_avg"] = float(rmse.mean())
    return row


def main():
    cli = parse_args()
    os.makedirs(os.path.dirname(cli.out), exist_ok=True)
    device = t.device(cli.device)
    run_args = dict(base_args)
    run_args["device"] = device
    run_args["train_flag"] = False

    rows = [
        eval_checkpoint(label, path, run_args, device, cli)
        for label, path in CHECKPOINTS.items()
    ]
    fieldnames = [
        "model",
        "checkpoint",
        "n_samples",
        "pos_rmse_0.5s",
        "pos_rmse_1.0s",
        "pos_rmse_1.5s",
        "pos_rmse_2.0s",
        "pos_rmse_avg",
    ]
    with open(cli.out, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    for row in rows:
        print(row)
    print(f"Saved to: {cli.out}")


if __name__ == "__main__":
    main()