Skip to content

Route console messages through tqdm.write() to avoid corrupted progress bars #1606

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dxqb wants to merge 1 commit into
base: master
Choose a base branch
from
+16 −14
Open

Route console messages through tqdm.write() to avoid corrupted progress bars #1606

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions modules/trainer/GenericTrainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,7 @@ def __prune_backups(self, backups_to_keep: int):
try:
shutil.rmtree(dirpath)
except Exception:
print(f"Could not delete old rolling backup {dirpath}")
tqdm.write(f"Could not delete old rolling backup {dirpath}")

return

Expand Down Expand Up @@ -270,7 +270,7 @@ def on_sample_custom(sampler_output: ModelSamplerOutput):
)
except Exception:
traceback.print_exc()
print("Error during sampling, proceeding without sampling")
tqdm.write("Error during sampling, proceeding without sampling")

torch_gc()

Expand Down Expand Up @@ -303,7 +303,7 @@ def __sample_during_training(
# We absolutely do not want to fail training just because the sample definition file becomes missing or broken right before sampling.
except Exception:
traceback.print_exc()
print("Error during loading the sample definition file, proceeding without sampling")
tqdm.write("Error during loading the sample definition file, proceeding without sampling")
sample_params_list = []

if self.model.ema:
Expand Down Expand Up @@ -428,7 +428,7 @@ def __save_backup_config(self, backup_path):
if os.path.isfile(self.config.sample_definition_file_name):
shutil.copy2(self.config.sample_definition_file_name, samples_path)

def __backup(self, train_progress: TrainProgress, print_msg: bool = True, print_cb: Callable[[str], None] = print):
def __backup(self, train_progress: TrainProgress, print_msg: bool = True):
torch_gc()

self.callbacks.on_update_status("Creating backup")
Expand All @@ -443,7 +443,7 @@ def __backup(self, train_progress: TrainProgress, print_msg: bool = True, print_

try:
if print_msg:
print_cb("Creating Backup " + backup_path)
tqdm.write("Creating Backup " + backup_path)

self.model_saver.save(
self.model,
Expand All @@ -456,13 +456,13 @@ def __backup(self, train_progress: TrainProgress, print_msg: bool = True, print_
self.__save_backup_config(backup_path)
except Exception:
traceback.print_exc()
print("Could not save backup. Check your disk space!")
tqdm.write("Could not save backup. Check your disk space!")
try:
if os.path.isdir(backup_path):
shutil.rmtree(backup_path)
except Exception:
traceback.print_exc()
print("Could not delete partial backup")
tqdm.write("Could not delete partial backup")
finally:
if self.config.rolling_backup:
self.__prune_backups(self.config.rolling_backup_count)
Expand All @@ -475,7 +475,7 @@ def __backup(self, train_progress: TrainProgress, print_msg: bool = True, print_

torch_gc()

def __save(self, train_progress: TrainProgress, print_msg: bool = True, print_cb: Callable[[str], None] = print):
def __save(self, train_progress: TrainProgress, print_msg: bool = True):
torch_gc()

self.callbacks.on_update_status("Saving")
Expand All @@ -486,7 +486,7 @@ def __save(self, train_progress: TrainProgress, print_msg: bool = True, print_cb
f"{self.config.save_filename_prefix}{get_string_timestamp()}-save-{train_progress.filename_string()}{self.config.output_model_format.file_extension()}"
)
if print_msg:
print_cb("Saving " + save_path)
tqdm.write("Saving " + save_path)

try:
if self.model.ema:
Expand All @@ -508,13 +508,13 @@ def __save(self, train_progress: TrainProgress, print_msg: bool = True, print_cb
self.model.optimizer.train()
except Exception:
traceback.print_exc()
print("Could not save model. Check your disk space!")
tqdm.write("Could not save model. Check your disk space!")
try:
if os.path.isfile(save_path):
shutil.rmtree(save_path)
except Exception:
traceback.print_exc()
print("Could not delete partial save")
tqdm.write("Could not delete partial save")
finally:
if self.model.ema:
self.model.ema.copy_temp_to(self.parameters)
Expand Down Expand Up @@ -719,9 +719,9 @@ def sample_commands_fun():
if multi.is_master() and (backup or save):
self.model.to(self.temp_device)
if backup:
self.__backup(train_progress, True, step_tqdm.write)
self.__backup(train_progress, True)
if save:
self.__save(train_progress, True, step_tqdm.write)
self.__save(train_progress, True)
self.model_setup.setup_train_device(self.model, self.config)

self.callbacks.on_update_status("Training ...")
Expand Down
4 changes: 3 additions & 1 deletion modules/util/multi_gpu_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@

import torch

from tqdm import tqdm


def is_enabled() -> bool:
return torch.distributed.is_available() and torch.distributed.is_initialized()
Expand Down Expand Up @@ -149,7 +151,7 @@ def parameter_divergence(params: list[torch.Tensor], train_device: torch.device)
def warn_parameter_divergence(params: list[torch.Tensor], train_device: torch.device):
divergence = parameter_divergence(params, train_device)
if divergence is not None and divergence > 0:
print(f"\n\nWARNING: Parameter divergence between GPUs of {divergence}\n\n")
tqdm.write(f"WARNING: Parameter divergence between GPUs of {divergence}")


@torch.no_grad()
Expand Down