70 lines
3.3 KiB
Python
70 lines
3.3 KiB
Python
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)
|