Skip to content

Commit 10279c1

Browse files
perf(load_ml): halve the memory of a training pass
Profiling a plan cycle showed ML training dominating peak memory at 1.4GB, with the feature matrix and its copies accounting for most of it. A curriculum pass over three weeks of history builds 11,500 samples by 1,446 float32 features - 66.7MB - and the training path held several copies of it at once. Four changes, measured through the production code path at that scale: Preallocate the feature matrix. _create_dataset built a Python list of 11,500 separate per-sample arrays and then copied it into np.array(), keeping both alive. Rows are now written straight into a preallocated matrix, which also removes the five temporary arrays and the concatenated result that _build_sample created per sample. Normalise in place. _normalize_features gains an opt-in in_place flag, used at the three training sites where the caller is provably finished with the array it hands in - only the normalised arrays are read from there on. The default still copies. Compute normalisation statistics over row blocks. np.std(X, axis=0) materialises the deviations array internally, another full copy of the matrix to produce 1,446 numbers. _feature_mean_std accumulates over blocks in float64 instead, which needs no temporary larger than one block and is also far more accurate: on a large-mean dataset the fitted std error drops from 2.4e-06 to 4.6e-10 and the mean error from 2.2e-04 to the float32 representation floor. Stop promoting the matrix to float64. _get_min_std_array returned a float64 array, so np.maximum made feature_std float64 and (X - mean) / std promoted the entire normalised matrix to double precision - 133.3MB rather than 66.6MB - and ran every forward and backward pass in float64 against float32 weights. The array is now float32, matching the dataset and the weights. Measured on 11,519 samples through _create_dataset plus _normalize_features: dataset construction 181.8MB -> 107.9MB peak growth 298.9MB -> 165.1MB (-45%) normalised matrix float64 133.3MB -> float32 66.6MB The dtype fix changes results slightly, so this is not a pure refactor: across the ML test suite val_mae moves by 0.0054 kWh on average, early epochs are consistently a little better and training runs further before early stopping. All 29 ML tests pass before and after. The feature matrix itself is unchanged - a checksum over it is pinned in the new tests and does not move. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent e86772e commit 10279c1

3 files changed

Lines changed: 362 additions & 33 deletions

File tree

‎apps/predbat/load_predictor.py‎

Lines changed: 83 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -611,14 +611,23 @@ def _create_dataset(self, load_minutes, now_utc, pv_minutes=None, temp_minutes=N
611611
# Validation window: most recent chunks (chunk_idx 0 … validation_end_chunk-1)
612612
validation_end_chunk = validation_holdout_hours * 60 // CHUNK_MINUTES
613613

614-
X_train_list = []
614+
# Preallocate the feature matrices and write each sample straight into its row. Holding
615+
# a Python list of per-sample arrays and then copying it into np.array() keeps both
616+
# alive at once, which for a three week window is an extra 70MB on top of the 66MB
617+
# matrix. Rows are sized for the maximum possible sample count and the filled prefix is
618+
# returned; samples are only skipped where the history has gaps, so the slack is small.
619+
max_train_rows = max(max_chunk_idx - LOOKBACK_STEPS, 0)
620+
max_val_rows = max(validation_end_chunk, 0)
621+
X_train_all = np.empty((max_train_rows, TOTAL_FEATURES), dtype=np.float32)
622+
X_val_all = np.empty((max_val_rows, TOTAL_FEATURES), dtype=np.float32)
623+
train_rows = 0
624+
val_rows = 0
615625
y_train_list = []
616626
weight_list = []
617-
X_val_list = []
618627
y_val_list = []
619628

620-
def _build_sample(target_chunk_idx):
621-
"""Build one (features, target) sample centred on target_chunk_idx."""
629+
def _build_sample(target_chunk_idx, out_row):
630+
"""Write one sample's features into out_row and return its target, or None on a gap."""
622631
lookback_start = target_chunk_idx + 1
623632
lookback_values = []
624633
pv_lookback_values = []
@@ -651,25 +660,22 @@ def _build_sample(target_chunk_idx):
651660
day_of_year = target_time.timetuple().tm_yday
652661
time_features = self._create_time_features(minute_of_day, day_of_week, day_of_year)
653662

654-
features = np.concatenate(
655-
[
656-
np.array(lookback_values, dtype=np.float32),
657-
np.array(pv_lookback_values, dtype=np.float32),
658-
np.array(temp_lookback_values, dtype=np.float32),
659-
np.array(import_rate_lookback, dtype=np.float32),
660-
np.array(export_rate_lookback, dtype=np.float32),
661-
time_features,
662-
]
663-
)
664-
return features, np.array([target_value], dtype=np.float32)
663+
# Assigning the lists straight into row slices converts them in place, avoiding the
664+
# five temporary arrays and the concatenated result this used to build per sample
665+
offset = 0
666+
for values in (lookback_values, pv_lookback_values, temp_lookback_values, import_rate_lookback, export_rate_lookback):
667+
out_row[offset : offset + LOOKBACK_STEPS] = values
668+
offset += LOOKBACK_STEPS
669+
out_row[offset:] = time_features
670+
return np.array([target_value], dtype=np.float32)
665671

666672
# Training samples: all available chunks
667673
for target_chunk_idx in range(0, max_chunk_idx - LOOKBACK_STEPS):
668-
features, target = _build_sample(target_chunk_idx)
669-
if features is None:
674+
target = _build_sample(target_chunk_idx, X_train_all[train_rows])
675+
if target is None:
670676
continue
671677

672-
X_train_list.append(features)
678+
train_rows += 1
673679
y_train_list.append(target)
674680

675681
# Time-decay weighting (older samples get lower weight)
@@ -678,23 +684,23 @@ def _build_sample(target_chunk_idx):
678684

679685
# Validation samples: most recent validation_end_chunk chunks
680686
for target_chunk_idx in range(0, validation_end_chunk):
681-
features, target = _build_sample(target_chunk_idx)
682-
if features is None:
687+
target = _build_sample(target_chunk_idx, X_val_all[val_rows])
688+
if target is None:
683689
continue
684-
X_val_list.append(features)
690+
val_rows += 1
685691
y_val_list.append(target)
686692

687-
if not X_train_list:
693+
if not train_rows:
688694
return None, None, None, None, None
689695

690-
X_train = np.array(X_train_list, dtype=np.float32)
696+
X_train = X_train_all[:train_rows]
691697
y_train = np.array(y_train_list, dtype=np.float32)
692698
train_weights = np.array(weight_list, dtype=np.float32)
693699

694700
# Normalize weights to sum to number of samples
695701
train_weights = train_weights * len(train_weights) / np.sum(train_weights)
696702

697-
X_val = np.array(X_val_list, dtype=np.float32) if X_val_list else None
703+
X_val = X_val_all[:val_rows] if val_rows else None
698704
y_val = np.array(y_val_list, dtype=np.float32) if y_val_list else None
699705

700706
return X_train, y_train, train_weights, X_val, y_val
@@ -808,6 +814,39 @@ def _ar_rollout_diagnostic(self, load_minutes, now_utc, pv_minutes=None, temp_mi
808814

809815
return float(np.mean(errors)), float(np.mean(biases))
810816

817+
@staticmethod
818+
def _feature_mean_std(X, block=4096):
819+
"""
820+
Compute per-feature mean and std over row blocks, accumulating in float64.
821+
822+
np.mean/np.std on a float32 array accumulate in float32, which loses low bits on the
823+
feature groups that sit far from zero, and np.std materialises the deviations array
824+
internally - another full copy of a 66MB training matrix to produce 1,446 numbers.
825+
Two passes over row blocks avoids both: no temporary larger than one block, and
826+
float64 accumulators that match a float64 reference.
827+
828+
Args:
829+
X: Feature array, shape (samples, features)
830+
block: Rows accumulated per pass
831+
832+
Returns:
833+
Tuple of (mean, std) as float32 arrays
834+
"""
835+
rows = len(X)
836+
total = np.zeros(X.shape[1], dtype=np.float64)
837+
for start in range(0, rows, block):
838+
total += X[start : start + block].sum(axis=0, dtype=np.float64)
839+
mean = total / rows
840+
841+
squares = np.zeros(X.shape[1], dtype=np.float64)
842+
for start in range(0, rows, block):
843+
deviation = X[start : start + block].astype(np.float64)
844+
deviation -= mean
845+
squares += np.square(deviation).sum(axis=0)
846+
std = np.sqrt(squares / rows)
847+
848+
return mean.astype(np.float32), std.astype(np.float32)
849+
811850
def _get_min_std_array(self, n_features):
812851
"""
813852
Return the per-feature minimum std array used to prevent extreme normalization.
@@ -818,7 +857,10 @@ def _get_min_std_array(self, n_features):
818857
Returns:
819858
numpy array of minimum std values, shape (n_features,)
820859
"""
821-
min_std = np.ones(n_features) * 1e-8 # Default fallback
860+
# float32 to match the dataset and weights: a float64 minimum promotes the whole
861+
# normalised feature matrix to float64 via np.maximum, doubling its size and running
862+
# training in double precision against float32 weights
863+
min_std = np.ones(n_features, dtype=np.float32) * 1e-8 # Default fallback
822864
if n_features == TOTAL_FEATURES:
823865
min_std[0:LOOKBACK_STEPS] = 0.01 # Load energy (kWh)
824866
min_std[LOOKBACK_STEPS : 2 * LOOKBACK_STEPS] = 0.01 # PV energy (kWh)
@@ -858,7 +900,7 @@ def _log_normalization_stats(self, label=""):
858900

859901
self.log("ML Predictor: Normalization stats [{}] target(mean={:.4f} std={:.4f}) {}".format(label, self.target_mean if self.target_mean is not None else 0, self.target_std if self.target_std is not None else 0, " ".join(parts)))
860902

861-
def _normalize_features(self, X, fit=False, ema_alpha=0.0):
903+
def _normalize_features(self, X, fit=False, ema_alpha=0.0, in_place=False):
862904
"""
863905
Normalize features using z-score normalization with feature-specific minimum stds.
864906
@@ -868,22 +910,25 @@ def _normalize_features(self, X, fit=False, ema_alpha=0.0):
868910
ema_alpha: If > 0 and existing params exist, blend new stats with old via EMA
869911
(new = alpha * new_stats + (1-alpha) * old_stats). Used during
870912
fine-tuning to track feature distribution drift without sudden jumps.
913+
in_place: If True, write the normalised values back into X instead of building a
914+
new array. A training feature matrix is around 66MB, and the caller
915+
keeps no use for the un-normalised values, so copying doubles the
916+
resident cost of a training pass for nothing. Only pass this when the
917+
caller is finished with the array it hands in.
871918
872919
Returns:
873920
Normalized feature array
874921
"""
875922
if fit:
876-
self.feature_mean = np.mean(X, axis=0)
877-
self.feature_std = np.std(X, axis=0)
923+
self.feature_mean, self.feature_std = self._feature_mean_std(X)
878924

879925
# Clamp std to per-feature minimums to prevent extreme normalization
880926
self.feature_std = np.maximum(self.feature_std, self._get_min_std_array(len(self.feature_std)))
881927
self._log_normalization_stats(label="fit")
882928

883929
elif ema_alpha > 0 and self.feature_mean is not None and self.feature_std is not None:
884930
# EMA update: blend new statistics with existing to track distribution drift
885-
new_mean = np.mean(X, axis=0)
886-
new_std = np.std(X, axis=0)
931+
new_mean, new_std = self._feature_mean_std(X)
887932

888933
# Apply same min-std clamping to new stats before blending
889934
new_std = np.maximum(new_std, self._get_min_std_array(len(new_std)))
@@ -896,6 +941,11 @@ def _normalize_features(self, X, fit=False, ema_alpha=0.0):
896941
if self.feature_mean is None or self.feature_std is None:
897942
return X
898943

944+
if in_place:
945+
X -= self.feature_mean
946+
X /= self.feature_std
947+
return X
948+
899949
return (X - self.feature_mean) / self.feature_std
900950

901951
def _normalize_targets(self, y, fit=False):
@@ -1039,13 +1089,13 @@ def train(
10391089
# On initial train: fit normalization from scratch
10401090
# On fine-tune: apply EMA update to track distribution drift gradually
10411091
if is_initial or not self.model_initialized:
1042-
X_train_norm = self._normalize_features(X_train, fit=True)
1092+
X_train_norm = self._normalize_features(X_train, fit=True, in_place=True)
10431093
y_train_norm = self._normalize_targets(y_train, fit=True)
10441094
else:
1045-
X_train_norm = self._normalize_features(X_train, fit=False, ema_alpha=norm_ema_alpha)
1095+
X_train_norm = self._normalize_features(X_train, fit=False, ema_alpha=norm_ema_alpha, in_place=True)
10461096
y_train_norm = self._normalize_targets(y_train, fit=False)
10471097
self.log("ML Predictor: Applied EMA normalization update (alpha={}) to track feature drift".format(norm_ema_alpha))
1048-
X_val_norm = self._normalize_features(X_val, fit=False)
1098+
X_val_norm = self._normalize_features(X_val, fit=False, in_place=True)
10491099
y_val_norm = self._normalize_targets(y_val, fit=False)
10501100

10511101
# Initialise weights if needed

0 commit comments

Comments
 (0)