Buffer predict rows across row-group boundaries before inference
Inference batches were capped at the source parquet's row-group size (e.g. 122,880 rows) because the old loop only sliced within a single chunk read from disk. --batch-size (including auto) had no effect once it exceeded that, leaving most estimated GPU memory unused. Accumulate rows across row groups and files into a buffer and slice exactly batch-size pieces off it, so inference always uses the requested batch size regardless of how the file happens to be chunked.
This commit is contained in:
+122
-110
@@ -350,6 +350,118 @@ def predict(
|
||||
total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
|
||||
chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks
|
||||
|
||||
def _concat(
|
||||
a: dict[str, np.ndarray], b: dict[str, np.ndarray]
|
||||
) -> dict[str, np.ndarray]:
|
||||
return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}
|
||||
|
||||
def _process(piece: dict[str, np.ndarray]) -> None:
|
||||
nonlocal writer, total
|
||||
|
||||
if coord == Coord.local:
|
||||
cond_cont, cond_cat, target_raw, _, _ = build_features(
|
||||
piece, pdg_map, mat_map
|
||||
)
|
||||
cond_cont = cond_norm.transform(cond_cont)
|
||||
else:
|
||||
cond_cont, cond_cat = build_cond_features(
|
||||
piece, pdg_map, mat_map, cond_norm
|
||||
)
|
||||
|
||||
cc = torch.from_numpy(cond_cont).float().to(_device)
|
||||
ck = torch.from_numpy(cond_cat).long().to(_device)
|
||||
pred = sample_flow(model, cc, ck, steps=steps).cpu().numpy() # normalised
|
||||
|
||||
# Inverse-normalise → local frame, log-scaled scalars
|
||||
raw = tgt_norm.inverse_transform(pred)
|
||||
|
||||
if coord == Coord.local:
|
||||
table = pa.table(
|
||||
{
|
||||
"event_id": piece["event_id"],
|
||||
"pdg": piece["pdg"],
|
||||
"pre_x": piece["pre_pos"][:, 0],
|
||||
"pre_y": piece["pre_pos"][:, 1],
|
||||
"pre_z": piece["pre_pos"][:, 2],
|
||||
"pre_E": piece["pre_E"],
|
||||
"pre_dx": piece["pre_dir"][:, 0],
|
||||
"pre_dy": piece["pre_dir"][:, 1],
|
||||
"pre_dz": piece["pre_dir"][:, 2],
|
||||
"material": piece["material"],
|
||||
"layer_id": piece["layer_id"],
|
||||
"n_sec": piece["n_sec"],
|
||||
**{
|
||||
f"pred_{name}": raw[:, j]
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||||
},
|
||||
**{
|
||||
f"true_{name}": target_raw[:, j]
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||||
},
|
||||
}
|
||||
)
|
||||
else:
|
||||
step_length = inv_log_transform(raw[:, 0])
|
||||
delta_e = inv_log_transform(raw[:, 1])
|
||||
edep = inv_log_transform(raw[:, 2])
|
||||
|
||||
# Normalise predicted direction then rotate back to world frame
|
||||
post_dir_local = raw[:, 3:6].copy()
|
||||
norms = np.linalg.norm(post_dir_local, axis=1, keepdims=True)
|
||||
post_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
||||
post_dir_world = inv_local_frame_rotation(piece["pre_dir"], post_dir_local)
|
||||
|
||||
# Same for the travel direction, then reconstruct post_pos from
|
||||
# the single shared step_length so the two stay consistent.
|
||||
travel_dir_local = raw[:, 6:9].copy()
|
||||
norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
|
||||
travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
||||
post_pos_world = reconstruct_post_pos(
|
||||
piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
|
||||
)
|
||||
|
||||
table = pa.table(
|
||||
{
|
||||
"event_id": piece["event_id"],
|
||||
"pdg": piece["pdg"],
|
||||
"pre_x": piece["pre_pos"][:, 0],
|
||||
"pre_y": piece["pre_pos"][:, 1],
|
||||
"pre_z": piece["pre_pos"][:, 2],
|
||||
"pre_E": piece["pre_E"],
|
||||
"pre_dx": piece["pre_dir"][:, 0],
|
||||
"pre_dy": piece["pre_dir"][:, 1],
|
||||
"pre_dz": piece["pre_dir"][:, 2],
|
||||
"material": piece["material"],
|
||||
"layer_id": piece["layer_id"],
|
||||
"n_sec": piece["n_sec"],
|
||||
"step_length": step_length,
|
||||
"delta_e": delta_e,
|
||||
"edep": edep,
|
||||
"post_dx": post_dir_world[:, 0],
|
||||
"post_dy": post_dir_world[:, 1],
|
||||
"post_dz": post_dir_world[:, 2],
|
||||
"post_x": post_pos_world[:, 0],
|
||||
"post_y": post_pos_world[:, 1],
|
||||
"post_z": post_pos_world[:, 2],
|
||||
}
|
||||
)
|
||||
|
||||
table = table.replace_schema_metadata(
|
||||
{
|
||||
PREDICT_COORD_METADATA_KEY: coord.value,
|
||||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||||
}
|
||||
)
|
||||
|
||||
if writer is None:
|
||||
writer = pq.ParquetWriter(out, table.schema)
|
||||
writer.write_table(table)
|
||||
total += len(piece["event_id"])
|
||||
|
||||
# Buffer rows across row-group boundaries so the inference batch size
|
||||
# isn't capped by however the source file happens to be chunked.
|
||||
buffer: dict[str, np.ndarray] | None = None
|
||||
|
||||
bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
|
||||
for path in files:
|
||||
for chunk in chunk_iter(path):
|
||||
@@ -360,119 +472,19 @@ def predict(
|
||||
unknown_pdg_counts.update(int(p) for p in chunk["pdg"][~pdg_mask])
|
||||
chunk = {k: v[pdg_mask] for k, v in chunk.items()}
|
||||
|
||||
N = len(chunk["event_id"])
|
||||
skipped += N_in - N
|
||||
if N == 0:
|
||||
bar.update(N_in)
|
||||
skipped += N_in - len(chunk["event_id"])
|
||||
bar.update(N_in)
|
||||
if len(chunk["event_id"]) == 0:
|
||||
continue
|
||||
|
||||
if coord == Coord.local:
|
||||
cond_cont, cond_cat, target_raw, _, _ = build_features(
|
||||
chunk, pdg_map, mat_map
|
||||
)
|
||||
cond_cont = cond_norm.transform(cond_cont)
|
||||
else:
|
||||
cond_cont, cond_cat = build_cond_features(
|
||||
chunk, pdg_map, mat_map, cond_norm
|
||||
)
|
||||
buffer = chunk if buffer is None else _concat(buffer, chunk)
|
||||
while len(buffer["event_id"]) >= bs:
|
||||
piece = {k: v[:bs] for k, v in buffer.items()}
|
||||
buffer = {k: v[bs:] for k, v in buffer.items()}
|
||||
_process(piece)
|
||||
|
||||
# Inference in batch_size slices
|
||||
pred_parts = []
|
||||
for start in range(0, N, bs):
|
||||
end = min(start + bs, N)
|
||||
cc = torch.from_numpy(cond_cont[start:end]).float().to(_device)
|
||||
ck = torch.from_numpy(cond_cat[start:end]).long().to(_device)
|
||||
pred_parts.append(sample_flow(model, cc, ck, steps=steps).cpu().numpy())
|
||||
pred = np.concatenate(pred_parts, axis=0) # (N, 9) normalised
|
||||
|
||||
# Inverse-normalise → local frame, log-scaled scalars
|
||||
raw = tgt_norm.inverse_transform(pred)
|
||||
|
||||
if coord == Coord.local:
|
||||
table = pa.table(
|
||||
{
|
||||
"event_id": chunk["event_id"],
|
||||
"pdg": chunk["pdg"],
|
||||
"pre_x": chunk["pre_pos"][:, 0],
|
||||
"pre_y": chunk["pre_pos"][:, 1],
|
||||
"pre_z": chunk["pre_pos"][:, 2],
|
||||
"pre_E": chunk["pre_E"],
|
||||
"pre_dx": chunk["pre_dir"][:, 0],
|
||||
"pre_dy": chunk["pre_dir"][:, 1],
|
||||
"pre_dz": chunk["pre_dir"][:, 2],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
**{
|
||||
f"pred_{name}": raw[:, j]
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||||
},
|
||||
**{
|
||||
f"true_{name}": target_raw[:, j]
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||||
},
|
||||
}
|
||||
)
|
||||
else:
|
||||
step_length = inv_log_transform(raw[:, 0])
|
||||
delta_e = inv_log_transform(raw[:, 1])
|
||||
edep = inv_log_transform(raw[:, 2])
|
||||
|
||||
# Normalise predicted direction then rotate back to world frame
|
||||
post_dir_local = raw[:, 3:6].copy()
|
||||
norms = np.linalg.norm(post_dir_local, axis=1, keepdims=True)
|
||||
post_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
||||
post_dir_world = inv_local_frame_rotation(
|
||||
chunk["pre_dir"], post_dir_local
|
||||
)
|
||||
|
||||
# Same for the travel direction, then reconstruct post_pos from
|
||||
# the single shared step_length so the two stay consistent.
|
||||
travel_dir_local = raw[:, 6:9].copy()
|
||||
norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True)
|
||||
travel_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
||||
post_pos_world = reconstruct_post_pos(
|
||||
chunk["pre_pos"], chunk["pre_dir"], step_length, travel_dir_local
|
||||
)
|
||||
|
||||
table = pa.table(
|
||||
{
|
||||
"event_id": chunk["event_id"],
|
||||
"pdg": chunk["pdg"],
|
||||
"pre_x": chunk["pre_pos"][:, 0],
|
||||
"pre_y": chunk["pre_pos"][:, 1],
|
||||
"pre_z": chunk["pre_pos"][:, 2],
|
||||
"pre_E": chunk["pre_E"],
|
||||
"pre_dx": chunk["pre_dir"][:, 0],
|
||||
"pre_dy": chunk["pre_dir"][:, 1],
|
||||
"pre_dz": chunk["pre_dir"][:, 2],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
"step_length": step_length,
|
||||
"delta_e": delta_e,
|
||||
"edep": edep,
|
||||
"post_dx": post_dir_world[:, 0],
|
||||
"post_dy": post_dir_world[:, 1],
|
||||
"post_dz": post_dir_world[:, 2],
|
||||
"post_x": post_pos_world[:, 0],
|
||||
"post_y": post_pos_world[:, 1],
|
||||
"post_z": post_pos_world[:, 2],
|
||||
}
|
||||
)
|
||||
|
||||
table = table.replace_schema_metadata(
|
||||
{
|
||||
PREDICT_COORD_METADATA_KEY: coord.value,
|
||||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||||
}
|
||||
)
|
||||
|
||||
if writer is None:
|
||||
writer = pq.ParquetWriter(out, table.schema)
|
||||
writer.write_table(table)
|
||||
total += N
|
||||
bar.update(N_in)
|
||||
if buffer is not None and len(buffer["event_id"]) > 0:
|
||||
_process(buffer)
|
||||
|
||||
bar.close()
|
||||
if writer is not None:
|
||||
|
||||
Reference in New Issue
Block a user