Basic logging feature addition for Wandb + basic code

[kukaiN]

I use Wandb for logging but wandb makes random names (ex: dainty-haze-10, blooming-salad-13, ...) for the run and sometimes it's annoying finding which model is associated to a run. I attached below a simple addition to the wandb section in the code which upload the basic arguments and the output name so it's easier to navigate and compare runs on wandb.

add this to line 755~757 in sd_scripts\train_network.py in the kohya backend

# basically a dict of arguments passed to script
train_config = {
                "base_model": sd_model_name,
                "output_name": args.output_name,
                "learning_rate": args.learning_rate,
                "text_encoder_lr": args.text_encoder_lr,
                "unet_lr": args.unet_lr,
                "num_train_images": train_dataset_group.num_train_images,
                "num_reg_images": train_dataset_group.num_reg_images,
                "num_batches_per_epoch": len(train_dataloader),
                "num_epochs": num_train_epochs,
                "gradient_checkpointing": args.gradient_checkpointing,
                "gradient_accumulation_steps": args.gradient_accumulation_steps,
                "max_train_steps": args.max_train_steps,
                "lr_warmup_steps": args.lr_warmup_steps,
                "lr_scheduler": args.lr_scheduler,
                "network_module": args.network_module,
                "network_dim": args.network_dim,  # None means default because another network than LoRA may have another default dim
                "network_alpha": args.network_alpha,  # some networks may not have alpha
                "base_model_version": model_version,
                "clip_skip": args.clip_skip,
                "max_token_length": args.max_token_length,
                "noise_offset": args.noise_offset,
                "multires_noise_iterations": args.multires_noise_iterations,
                "multires_noise_discount": args.multires_noise_discount,
                "adaptive_noise_scale": args.adaptive_noise_scale,
                "zero_terminal_snr": args.zero_terminal_snr,
                "optimizer": optimizer_name + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),
                "min_snr_gamma": args.min_snr_gamma,
                "debiased_estimation": bool(args.debiased_estimation_loss),
                "noise_offset_random_strength": args.noise_offset_random_strength,
                "ip_noise_gamma_random_strength": args.ip_noise_gamma_random_strength,
                "loss_type": args.loss_type,
                "huber_schedule": args.huber_schedule,
                "huber_c": args.huber_c,
                }

if args.network_args: # this adds the extra network arg for special loras (locon, lyco, etc)
    for i,n_arg in enumerate(args.network_args):
        if "=" in n_arg: # if an "=" is in the string
            split_arg = n_arg.split("=", 1)
            train_config[split_arg[0]] = split_arg[1]
        else: # exception thingy, just in case there's an arg with no equal sign
            train_config["network_arg_"+str(i)] = n_arg

# added config=train_config, which sends the set values to wandb when the script starts
accelerator.init_trackers(
                "network_train" if args.log_tracker_name is None else args.log_tracker_name, config=train_config, init_kwargs=init_kwargs
            )

Here's what you'll see on Wandb, the run on the first line have the added code. Previous runs without the code will have blank entries:

[Jelosus2]

Sorry but we don't modify the sd-scripts code directly because is not something that we have developed. I'm not sure if Derrian has other way to accomplish this or if he sees this feature worth implementing. Thanks for your share by the way.

[Jelosus2]

Yea, it would be time consuming and annoying to maintain so we won't be implementing this. What you could do is open a Pull Request on the sd-scripts repository and if they merge it we could update our submodule to match that version. With that said I'll be closing this issue.