| 2 | from torch.utils.tensorboard import SummaryWriter |
| 3 | |
| 4 | class TensorboardCallback(Callback): |
| 5 | |
| 6 | def __init__(self, experiment_name): |
| 7 | self.tensorboard_writer = SummaryWriter("tensorboard_logs/"+experiment_name) |
| 8 | self.experiment_name=experiment_name |
| 9 | |
| 10 | |
| 11 | def after_forward_pass(self, phase, loss=0, loss_pos=0, loss_dir=0, loss_curv=0, loss_kl=0, lr=0, z_deviation=None, z=None, z_no_eps=None, **kwargs): |
| 12 | |
| 13 | if phase.iter_nr%300==0 and phase.grad: |
| 14 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss', loss.item(), phase.iter_nr) |
| 15 | if loss_pos!=0: |
| 16 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_pos', loss_pos.item(), phase.iter_nr) |
| 17 | # if loss_l1!=0: |
| 18 | # self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_l1', loss_l1.item(), phase.iter_nr) |
| 19 | # if loss_l2!=0: |
| 20 | # self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_l2', loss_l2.item(), phase.iter_nr) |
| 21 | if loss_dir!=0: |
| 22 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_dir', loss_dir.item(), phase.iter_nr) |
| 23 | if loss_curv!=0: |
| 24 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_curv', loss_curv.item(), phase.iter_nr) |
| 25 | if loss_kl!=0: |
| 26 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_kl', loss_kl.item(), phase.iter_nr) |
| 27 | |
| 28 | if lr!=0: |
| 29 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/lr', lr, phase.iter_nr) |
| 30 | |
| 31 | if z_deviation is not None: |
| 32 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/z_deviation', z_deviation.std(), phase.iter_nr) |
| 33 | |
| 34 | # if z is not None: |
| 35 | # self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/z_max', z.max(), phase.iter_nr) |
| 36 | if z_no_eps is not None: |
| 37 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/z_no_eps_mean', z_no_eps.mean(), phase.iter_nr) |
| 38 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/z_no_eps_std', z_no_eps.std(), phase.iter_nr) |
| 39 | |
| 40 | |
| 41 | def epoch_ended(self, phase, **kwargs): |
| 42 | avg_loss_pos=phase.loss_pos_acum_per_epoch/phase.samples_processed_this_epoch |
| 43 | avg_loss_dir=phase.loss_dir_acum_per_epoch/phase.samples_processed_this_epoch |
| 44 | avg_loss_curv=phase.loss_curv_acum_per_epoch/phase.samples_processed_this_epoch |
| 45 | |
| 46 | if phase.grad==False and phase.loss_pos_acum_per_epoch!=0: |
| 47 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_pos_avg', avg_loss_pos.item(), phase.epoch_nr) |
| 48 | if phase.grad==False and phase.loss_dir_acum_per_epoch!=0: |
| 49 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_dir_avg', avg_loss_dir.item(), phase.epoch_nr) |
| 50 | if phase.grad==False and phase.loss_curv_acum_per_epoch!=0: |
| 51 | self.tensorboard_writer.add_scalar('hair_forge/' + phase.name + '/loss_curv_avg', avg_loss_curv.item(), phase.epoch_nr) |
| 52 | |