From ad02e9bdac7c988eec81ee06be862763b01a03f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicolas=20B=C3=A9reux?= Date: Wed, 24 Jun 2026 13:21:36 +0200 Subject: [PATCH] save effective time --- rbms/io.py | 2 ++ rbms/training/pcd.py | 6 +++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/rbms/io.py b/rbms/io.py index 24a4b1c..6798f9a 100644 --- a/rbms/io.py +++ b/rbms/io.py @@ -16,6 +16,7 @@ def save_model( num_updates: int, time: float, learning_rate: Tensor, + effective_time: np.ndarray, flags: list[str] = [], ) -> None: """Save the current state of the model. @@ -48,6 +49,7 @@ def save_model( checkpoint["numpy_rng_arg3"] = np.random.get_state()[3] checkpoint["numpy_rng_arg4"] = np.random.get_state()[4] checkpoint["time"] = time + checkpoint["effective_time"] = effective_time checkpoint["learning_rate"] = learning_rate.cpu().numpy() # Update the parallel chains to resume training if "parallel_chains" in f.keys(): diff --git a/rbms/training/pcd.py b/rbms/training/pcd.py index e1c2bf9..4eae5b0 100644 --- a/rbms/training/pcd.py +++ b/rbms/training/pcd.py @@ -37,7 +37,7 @@ def train( pbar.set_description(f"Training {params.name}") start = time.perf_counter() - + effective_time = np.zeros(len(optimizer)) for idx in range(curr_update + 1, num_updates + 1): batch = train_dataset.batch(batch_size) data, weights = batch["data"], batch["weights"] @@ -69,7 +69,8 @@ def train( params.post_grad_update() sampler.post_grad_update(params=params) - + learning_rates = np.asarray([opt.param_groups[0]["lr"] for opt in optimizer]) + effective_time += learning_rates # Get flags for save flags = [] flags = params.save_flags(flags) @@ -81,7 +82,6 @@ def train( names_params = ( list(params.named_parameters().keys()) if len(optimizer) > 1 else ["all"] ) - learning_rates = np.asarray([opt.param_groups[0]["lr"] for opt in optimizer]) metrics = {} metrics = sampler.get_metrics_display(