Spaces:
Runtime error
Runtime error
| import os | |
| import numpy as np | |
| import torch | |
| from torch.utils.tensorboard import SummaryWriter | |
| from modules import shared | |
| def tensorboard_setup(log_directory): | |
| os.makedirs(os.path.join(log_directory, "tensorboard"), exist_ok=True) | |
| return SummaryWriter( | |
| log_dir=os.path.join(log_directory, "tensorboard"), | |
| flush_secs=shared.opts.training_tensorboard_flush_every) | |
| def tensorboard_log_hyperparameter(tensorboard_writer:SummaryWriter, **kwargs): | |
| for keys in kwargs: | |
| if type(kwargs[keys]) not in [bool, str, float, int,None]: | |
| kwargs[keys] = str(kwargs[keys]) | |
| tensorboard_writer.add_hparams({ | |
| 'lr' : kwargs.get('lr', 0.01), | |
| 'GA steps' : kwargs.get('GA_steps', 1), | |
| 'bsize' : kwargs.get('batch_size', 1), | |
| 'layer structure' : kwargs.get('layer_structure', '1,2,1'), | |
| 'activation' : kwargs.get('activation', 'Linear'), | |
| 'weight_init' : kwargs.get('weight_init', 'Normal'), | |
| 'dropout_structure' : kwargs.get('dropout_structure', '0,0,0'), | |
| 'steps' : kwargs.get('max_steps', 10000), | |
| 'latent sampling': kwargs.get('latent_sampling_method', 'once'), | |
| 'template file': kwargs.get('template', 'nothing'), | |
| 'CosineAnnealing' : kwargs.get('CosineAnnealing', False), | |
| 'beta_repeat epoch': kwargs.get('beta_repeat_epoch', 0), | |
| 'epoch_mult':kwargs.get('epoch_mult', 1), | |
| 'warmup_step' : kwargs.get('warmup', 5), | |
| 'min_lr' : kwargs.get('min_lr', 6e-7), | |
| 'decay' : kwargs.get('gamma_rate', 1), | |
| 'adamW' : kwargs.get('adamW_opts', False), | |
| 'adamW_decay' : kwargs.get('adamW_decay', 0.01), | |
| 'adamW_beta1' : kwargs.get('adamW_beta_1', 0.9), | |
| 'adamW_beta2': kwargs.get('adamW_beta_2', 0.99), | |
| 'adamW_eps': kwargs.get('adamW_eps', 1e-8), | |
| 'gradient_clip_opt':kwargs.get('gradient_clip', 'None'), | |
| 'gradient_clip_value' : kwargs.get('gradient_clip_value', 1e-1), | |
| 'gradient_clip_norm' : kwargs.get('gradient_clip_norm_type', 2) | |
| }, | |
| {'hparam/loss' : kwargs.get('loss', 0.0)} | |
| ) | |
| def tensorboard_add(tensorboard_writer:SummaryWriter, loss, global_step, step, learn_rate, epoch_num, base_name=""): | |
| prefix = base_name + "/" if base_name else "" | |
| tensorboard_add_scaler(tensorboard_writer, prefix+"Loss/train", loss, global_step) | |
| tensorboard_add_scaler(tensorboard_writer, prefix+f"Loss/train/epoch-{epoch_num}", loss, step) | |
| tensorboard_add_scaler(tensorboard_writer, prefix+"Learn rate/train", learn_rate, global_step) | |
| tensorboard_add_scaler(tensorboard_writer, prefix+f"Learn rate/train/epoch-{epoch_num}", learn_rate, step) | |
| def tensorboard_add_scaler(tensorboard_writer:SummaryWriter, tag, value, step): | |
| tensorboard_writer.add_scalar(tag=tag, | |
| scalar_value=value, global_step=step) | |
| def tensorboard_add_image(tensorboard_writer:SummaryWriter, tag, pil_image, step, base_name=""): | |
| # Convert a pil image to a torch tensor | |
| prefix = base_name + "/" if base_name else "" | |
| img_tensor = torch.as_tensor(np.array(pil_image, copy=True)) | |
| img_tensor = img_tensor.view(pil_image.size[1], pil_image.size[0], | |
| len(pil_image.getbands())) | |
| img_tensor = img_tensor.permute((2, 0, 1)) | |
| tensorboard_writer.add_image(prefix+tag, img_tensor, global_step=step) | |