diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 61b1c41..a51c51e 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -284,8 +284,6 @@ def main(kwargs_dict: dict = {}): outfile_path = trainer.outfile_path train_data = np.genfromtxt(outfile_path, dtype=float, delimiter=",", names=True) total_train_time = train_data["epoch_duration"].sum() - fom = 1.0 / total_train_time - adiak_value("FOM", fom) if "total_optimizer_steps" in train_data.dtype.names: optimizer_steps = np.atleast_1d(train_data["total_optimizer_steps"]) total_optimizer_steps = int(optimizer_steps[-1]) @@ -296,15 +294,16 @@ def main(kwargs_dict: dict = {}): else: total_optimizer_steps = int(getattr(trainer, "total_optimizer_steps", 0)) adiak_value("total_optimizer_steps", total_optimizer_steps) - log.info( - f"FOM = {fom} (1 / total_train_time={total_train_time:.6f} seconds). " - f"This FOM is specific to problem_scale={config.problem_scale}, " - f"target_dice={config.target_dice}, seed={config.seed}, " - f"total_optimizer_steps={total_optimizer_steps}." - ) epochs = np.atleast_1d(train_data["epoch"]) total_epochs = int(epochs[-1]) if config.epochs == -1: + fom = 1.0 / total_train_time + adiak_value("FOM", fom) + log.info( + f"FOM = {fom} (1 / total_train_time={total_train_time:.6f} seconds). " + f"This FOM is specific to problem_scale={config.problem_scale}, " + f"target_dice={config.target_dice}, seed={config.seed}." + ) extra_msg = f"Trained to >= {config.target_dice} validation dice score in {total_train_time:.2f} seconds, {total_epochs} epochs, {total_optimizer_steps} optimizer steps." else: extra_msg = f"Completed in {total_train_time:.2f} seconds, {total_epochs} epochs, {total_optimizer_steps} optimizer steps."