161 lines
6.5 KiB
Python
161 lines
6.5 KiB
Python
import datetime
|
|
import io
|
|
import logging
|
|
|
|
from matplotlib import pyplot as plt
|
|
from trainer import ConsoleLogger
|
|
from trainer.logging.console_logger import tcolors
|
|
from trainer.utils.distributed import rank_zero_only
|
|
|
|
logger = logging.getLogger("trainer")
|
|
|
|
class MetricsLogger(ConsoleLogger):
|
|
def __init__(self, ):
|
|
super().__init__()
|
|
self.cur_epoch_start = None
|
|
self.cur_epoch_end = None
|
|
self.cur_epoch = None
|
|
self.epoch_durations = []
|
|
self.max_epoch = None
|
|
self.batch_steps = None
|
|
self.global_steps = None
|
|
self.loss = []
|
|
self.avg_loss = []
|
|
self.loss_mel_ce = []
|
|
self.avg_loss_mel_ce = []
|
|
self.loss_text_ce = []
|
|
self.avg_loss_text_ce = []
|
|
self.learning_rates = []
|
|
self.total_duration = None
|
|
self.estimated_duration = None
|
|
self.start = None
|
|
|
|
@staticmethod
|
|
def get_now():
|
|
return datetime.datetime.now()
|
|
@staticmethod
|
|
def format_duration(seconds):
|
|
if seconds:
|
|
# Calculate hours, minutes, and seconds
|
|
hours = seconds // 3600
|
|
minutes = (seconds % 3600) // 60
|
|
seconds = seconds % 60
|
|
return f"{int(hours)}:{int(minutes)}:{int(seconds)}"
|
|
else:
|
|
return None
|
|
|
|
@rank_zero_only
|
|
def print_epoch_start(self, epoch, max_epoch, output_path=None):
|
|
self.cur_epoch = epoch
|
|
self.cur_epoch_start = self.get_now()
|
|
self.max_epoch = max_epoch
|
|
super().print_epoch_start(epoch, max_epoch, output_path)
|
|
|
|
@rank_zero_only
|
|
def print_train_start(self):
|
|
super().print_train_start()
|
|
|
|
@rank_zero_only
|
|
def print_train_step(self, batch_steps, step, global_step, loss_dict, avg_loss_dict):
|
|
self.batch_steps = batch_steps
|
|
self.global_steps = global_step
|
|
|
|
for key, value in loss_dict.items():
|
|
if key == 'loss_mel_ce':
|
|
self.loss_mel_ce.append((global_step, value))
|
|
elif key == 'loss':
|
|
self.loss.append((global_step, value))
|
|
elif key == 'loss_text_ce':
|
|
self.loss_text_ce.append((global_step, value))
|
|
elif key == 'current_lr':
|
|
self.learning_rates.append((global_step, value))
|
|
super().print_train_step(batch_steps, step, global_step, loss_dict, avg_loss_dict)
|
|
|
|
# pylint: disable=unused-argument
|
|
@rank_zero_only
|
|
def print_train_epoch_end(self, global_step, epoch, epoch_time, print_dict):
|
|
super().print_epoch_end(global_step, epoch, epoch_time, print_dict)
|
|
if self.estimated_duration:
|
|
self.log_with_flush(
|
|
"\n{}{} >> ETA: {} {}".format(tcolors.UNDERLINE, tcolors.BOLD, self.format_duration(self.estimated_duration), tcolors.ENDC),
|
|
)
|
|
|
|
@rank_zero_only
|
|
def print_eval_start(self):
|
|
super().print_eval_start()
|
|
|
|
@rank_zero_only
|
|
def print_eval_step(self, step, loss_dict, avg_loss_dict):
|
|
super().print_eval_step(step, loss_dict, avg_loss_dict)
|
|
|
|
@rank_zero_only
|
|
def print_epoch_end(self, epoch, avg_loss_dict):
|
|
self.cur_epoch_end = self.get_now()
|
|
self.epoch_durations.append((self.cur_epoch_end - self.cur_epoch_start).total_seconds())
|
|
for key, value in avg_loss_dict.items():
|
|
if key == 'avg_loss_mel_ce':
|
|
self.avg_loss_mel_ce.append((epoch, value))
|
|
elif key == 'avg_loss':
|
|
self.avg_loss.append((epoch, value))
|
|
elif key == 'avg_loss_text_ce':
|
|
self.avg_loss_text_ce.append((epoch, value))
|
|
|
|
if len(self.epoch_durations) > 1 and self.max_epoch:
|
|
self.total_duration = sum(self.epoch_durations)
|
|
additional_data_points_needed = self.max_epoch - len(self.epoch_durations)
|
|
if additional_data_points_needed > 0:
|
|
weighted_avg_duration = sum((i + 1) * duration for i, duration in enumerate(self.epoch_durations)) / sum(range(1, len(self.epoch_durations) + 1))
|
|
self.estimated_duration = (weighted_avg_duration * additional_data_points_needed)
|
|
|
|
super().print_epoch_end(epoch, avg_loss_dict)
|
|
|
|
def plot_metrics(self, show_loss_mel_ce=False, show_loss_text_ce=False, show_loss=True):
|
|
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(24, 18))
|
|
|
|
ax2.set_title('Steps')
|
|
ax2.set_xlabel('Steps')
|
|
ax2.set_ylabel('Value')
|
|
|
|
ax1.set_title('Epochs')
|
|
ax1.set_xlabel('Epoch')
|
|
ax1.set_ylabel('Value')
|
|
|
|
ax3.set_title('Learning Rates')
|
|
ax3.set_xlabel('Steps')
|
|
ax3.set_ylabel('Value')
|
|
|
|
if self.loss and self.loss_text_ce and self.loss_mel_ce:
|
|
if show_loss:
|
|
ax2.plot(*zip(*self.loss), color='green', linestyle='-', label='Loss')
|
|
ax2.fill_between(*zip(*self.loss), color='lightgreen', alpha=0.1)
|
|
if show_loss_text_ce:
|
|
ax2.plot(*zip(*self.loss_text_ce), color='red', linestyle='-', label='Loss Text')
|
|
ax2.fill_between(*zip(*self.loss_text_ce), color='lightcoral', alpha=0.1)
|
|
if show_loss_mel_ce:
|
|
ax2.plot(*zip(*self.loss_mel_ce), color='blue', linestyle='-', label='Loss MEL')
|
|
ax2.fill_between(*zip(*self.loss_mel_ce), color='lightblue', alpha=0.1)
|
|
ax2.legend()
|
|
|
|
if self.learning_rates:
|
|
ax3.plot(*zip(*self.learning_rates), linestyle='-', color='blue', label='Learning Rate')
|
|
ax3.fill_between(*zip(*self.learning_rates), color='lightblue', alpha=0.1)
|
|
ax3.legend()
|
|
|
|
if self.avg_loss and self.avg_loss_text_ce and self.avg_loss_mel_ce:
|
|
if show_loss:
|
|
ax1.plot(*zip(*self.avg_loss), color='green', linestyle='-', label='Avg Loss')
|
|
ax1.fill_between(*zip(*self.avg_loss), color='lightgreen', alpha=0.1)
|
|
if show_loss_text_ce:
|
|
ax1.plot(*zip(*self.avg_loss_text_ce), color='red', linestyle='-', label='Avg Loss Text')
|
|
ax1.fill_between(*zip(*self.avg_loss_text_ce), color='lightcoral', alpha=0.1)
|
|
if show_loss_mel_ce:
|
|
ax1.plot(*zip(*self.avg_loss_mel_ce), color='blue', linestyle='-', label='Avg Loss MEL')
|
|
ax1.fill_between(*zip(*self.avg_loss_mel_ce), color='lightblue', alpha=0.1)
|
|
ax1.set_xlim(left=0)
|
|
ax1.legend()
|
|
|
|
plt.style.use('dark_background')
|
|
plt.tight_layout()
|
|
plt.savefig("training_metrics.png", format='png', bbox_inches='tight', pad_inches=0.1)
|
|
plt.close()
|
|
return "training_metrics.png" |