diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index fc5af6e..518156c 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -77,6 +77,7 @@ def __init__(self, model, config, device, log): self.criterion = None self.ce_class_weights = None self.global_step = 0 + self.total_optimizer_steps = 0 self.start_epoch = -1 self.ps = getattr(self.config, "_parallel_strategy", None) self.spatial_mesh = None # Spatial mesh for use w/ DistConv @@ -345,6 +346,8 @@ def cleanup_or_resume(self): "train_dice", "val_dice", "epoch_duration", + "optimizer_steps", + "total_optimizer_steps", ] if self.world_rank == 0 and self.start_epoch == 1: with open(self.outfile_path, "a", newline="") as outfile: @@ -648,6 +651,7 @@ def train(self): epoch_start_time = time.time() train_dice_total = 0 epoch_loss = 0 # Accumulator for per-batch losses + epoch_optimizer_steps = 0 minibatch_time_s = None minibatch_events = [] @@ -699,6 +703,7 @@ def train(self): begin_code_region("update_loss") pbar.update(batch_size) self.global_step += 1 + epoch_optimizer_steps += 1 # Stay on GPU epoch_loss += batch_loss end_code_region("update_loss") @@ -710,6 +715,7 @@ def train(self): # Calculate overall loss as average of per-batch loss overall_loss = epoch_loss.item() / len(self.train_loader) + self.total_optimizer_steps += epoch_optimizer_steps # # Evaluate model on validation set, update LR if necessary @@ -759,7 +765,7 @@ def train(self): # train_dice = float(train_dice_total.item() / len(self.train_loader)) self.log.info( - f" epoch {epoch} | train_loss={overall_loss:.6f} | val_loss={val_loss_avg:.6f} | train_dice_score {train_dice:.6f} | val_dice_score {val_score:.6f} | lr {self._current_learning_rate():.8f}" + f" epoch {epoch} | train_loss={overall_loss:.6f} | val_loss={val_loss_avg:.6f} | train_dice_score {train_dice:.6f} | val_dice_score {val_score:.6f} | lr {self._current_learning_rate():.8f} | optimizer_steps {epoch_optimizer_steps} | total_optimizer_steps {self.total_optimizer_steps}" ) self.log.debug(f" writing to csv at {self.outfile_path}") if self.world_rank == 0: @@ -774,13 +780,15 @@ def train(self): str(train_dice), str(val_score), str(epoch_duration), + str(epoch_optimizer_steps), + str(self.total_optimizer_steps), ] ) + "\n" ) outfile.flush() print( - f"Epoch {epoch} completed in {epoch_duration:.6f} seconds. Total train time so far: {time.time() - start:.6f} seconds. Median of minibatch times: {minibatch_time_s:.6f} seconds." + f"Epoch {epoch} completed in {epoch_duration:.6f} seconds. Total train time so far: {time.time() - start:.6f} seconds. Median of minibatch times: {minibatch_time_s:.6f} seconds. Optimizer steps this epoch: {epoch_optimizer_steps}. Total optimizer steps: {self.total_optimizer_steps}." ) # @@ -816,3 +824,4 @@ def train(self): f"Median of epoch minibatch time medians: {minibatch_time_s:.6f} seconds." ) adiak_value("final_epochs", completed_epochs) + adiak_value("total_optimizer_steps", self.total_optimizer_steps) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index f0223d1..61b1c41 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -286,19 +286,28 @@ def main(kwargs_dict: dict = {}): 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]) + elif "optimizer_steps" in train_data.dtype.names: + total_optimizer_steps = int( + np.atleast_1d(train_data["optimizer_steps"]).sum() + ) + 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"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: - extra_msg = f"Trained to >= {config.target_dice} validation dice score in {total_train_time:.2f} seconds, {total_epochs} epochs." + 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." - ) + extra_msg = f"Completed in {total_train_time:.2f} seconds, {total_epochs} epochs, {total_optimizer_steps} optimizer steps." log.info( f"Benchmark run at scale {config.problem_scale} complete. \n{extra_msg}"