IVON_visualization_tool/ls_step_server.py at main · RinRin-32/IVON_visualization_tool · 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
import h5py
from bokeh.plotting import curdoc
from bokeh.models import ColumnDataSource
from bokeh.layouts import column, row
import json
import sys
import argparse
import os
from skimage import measure
from bokeh.plotting import output_file, save
import numpy as np
from visualizer.ls_decisionboundary import LSBoundaryVisualizer


def extract_boundary_lines(xx, yy, zz):
    contours = measure.find_contours(zz, level=0.5)
    xs, ys = [], []
    for contour in contours:
        xs.append(xx[0, 0] + contour[:, 1] * (xx[0, -1] - xx[0, 0]) / zz.shape[1])
        ys.append(yy[0, 0] + contour[:, 0] * (yy[-1, 0] - yy[0, 0]) / zz.shape[0])
    return xs, ys

parser = argparse.ArgumentParser(description="Launch the Bokeh server with an HDF5 file, this plot is to display changes in model behavior over training step.")
parser.add_argument("--file", type=str, required=True, help="Path to the HDF5 file")
parser.add_argument("--output", type=str, required=False, help="If specified filename, while running on python not bokeh serve, the html will be saved in ./output")
parser.add_argument("--scale_factor", type=int, default=3, help="Scale plotting of influence exponentially, default set at 3")

args = parser.parse_args()

h5_file = args.file

if not h5_file.lower().endswith(".h5"):
    print(f"Error: The input file '{h5_file}' is not an HDF5 (.h5) file.")
    sys.exit(1)

if not os.path.isfile(h5_file):
    print(f"Error: The file '{h5_file}' does not exist.")
    sys.exit(1)

if args.output is not None:
    os.makedirs('./output', exist_ok=True)
    output_file(filename=f"./output/{args.output}.html", title="Static HTML file", mode="inline")

with h5py.File(h5_file, "r") as f:
    read = f["config"]["config_data"][()]
    config_json = read.decode("utf-8")
    config = json.loads(config_json)
    dataset = config.get("dataset")
    max_epoch = config.get("max_epochs")
    total_batches = config.get("total_batch")
    max_step = total_batches * max_epoch

    X_coord = np.array(f["coord/X_train"])
    y_train = np.array(f["coord/y_train"])

    all_epoch_noises = [f[f"scores/step_{epoch}"]["noise"][()] for epoch in range(max_step)]

    xx = [f[f"scores/step_{epoch}"]["decision_boundary"]["xx"][:] for epoch in range(max_step)]
    yy = [f[f"scores/step_{epoch}"]["decision_boundary"]["yy"][:] for epoch in range(max_step)]
    Z = [f[f"scores/step_{epoch}"]["decision_boundary"]["Z"][:] for epoch in range(max_step)]

colors = ["white", "white"]
marker = ["circle", "star"]

xs = []
ys = []
for step in range(max_step):
    xx_step = xx[step]
    yy_step = yy[step]
    zz_step = Z[step]

    boundary_x, boundary_y = extract_boundary_lines(xx_step, yy_step, zz_step)
    xs.append(boundary_x)
    ys.append(boundary_y)

scaled_alphas_list = []
scaled_sizes_list = []
alpha_min, alpha_max = 0.2, 1.0
size_min, size_max = 5, 50
scaling_factor = args.scale_factor  # Adjust to control exaggeration

for epoch_noises in all_epoch_noises:
    normed_values = (epoch_noises - np.min(epoch_noises)) / (np.max(epoch_noises) - np.min(epoch_noises) + 1e-8)

    # Apply exponential transformation to exaggerate differences
    exp_values = normed_values ** scaling_factor

    alpha_assignments = alpha_min + (alpha_max - alpha_min) * exp_values
    size_assignments = size_min + (size_max - size_min) * exp_values

    scaled_alphas_list.append(alpha_assignments.tolist())
    scaled_sizes_list.append(size_assignments.tolist())

shared_resource = ColumnDataSource(data={
    "epoch": list(range(max_step)),
    "xs": xs,
    "ys": ys,
    "size": scaled_sizes_list,
    "alpha": scaled_alphas_list,
})

shared_source = ColumnDataSource(data={
    "x": X_coord[:, 0],
    "y": X_coord[:, 1],
    "class": y_train,
    "color": ['white'] * len(y_train),
    "marker": [marker[cls] for cls in y_train],
    "size": scaled_sizes_list[0],
    "alpha": scaled_alphas_list[0]
})

boundary = LSBoundaryVisualizer(shared_source, shared_resource, max_step, colors, total_batches, mode='Step')

boundary_layout = column(boundary.get_layout(), sizing_mode="scale_both")

layout = row(boundary_layout)

curdoc().add_root(layout)

if args.output is not None:
    layout.sizing_mode = "scale_both"
    save(layout)