351 lines
15 KiB
Python
351 lines
15 KiB
Python
import datetime
|
|
import os
|
|
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
|
|
self.current_model_path = None
|
|
|
|
@staticmethod
|
|
def get_now():
|
|
return datetime.datetime.now()
|
|
@staticmethod
|
|
def format_duration(seconds):
|
|
if seconds:
|
|
# Calculate hours, minutes, and seconds
|
|
hours = int(seconds // 3600)
|
|
minutes = int((seconds % 3600) // 60)
|
|
seconds = int(seconds % 60)
|
|
return f"{hours:02d}h {minutes:02d}m {seconds:02d}s" # More readable format
|
|
else:
|
|
return "N/A" # Better than None for display
|
|
|
|
@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 update_model_path(self, path):
|
|
"""Set current model path for log reading"""
|
|
try:
|
|
# Find all XTTS_FT folders in the given path
|
|
training_folders = []
|
|
for folder in os.listdir(path):
|
|
full_path = os.path.join(path, folder)
|
|
if os.path.isdir(full_path) and folder.startswith("XTTS_FT"):
|
|
training_folders.append(full_path)
|
|
|
|
if not training_folders:
|
|
print(f"[DEBUG] No XTTS_FT folders found in {path}")
|
|
self.current_model_path = path
|
|
return
|
|
|
|
# Get the most recent folder
|
|
latest_folder = max(training_folders, key=os.path.getctime)
|
|
# print(f"[DEBUG] Found latest training folder: {latest_folder}")
|
|
self.current_model_path = latest_folder
|
|
self.global_steps = 0 # Reset global steps
|
|
|
|
except Exception as e:
|
|
print(f"[DEBUG] Error updating model path: {e}")
|
|
self.current_model_path = path
|
|
|
|
def _format_metrics_for_log(self):
|
|
"""Format metrics for our log file"""
|
|
log_text = "\n=== Training Metrics Update ===\n"
|
|
log_text += f"Time: {self.get_now()}\n"
|
|
|
|
if self.avg_loss:
|
|
log_text += f"Current Average Loss: {self.avg_loss[-1][1]}\n"
|
|
if self.avg_loss_mel_ce:
|
|
log_text += f"Current MEL CE Loss: {self.avg_loss_mel_ce[-1][1]}\n"
|
|
if self.avg_loss_text_ce:
|
|
log_text += f"Current Text CE Loss: {self.avg_loss_text_ce[-1][1]}\n"
|
|
|
|
log_text += "===========================\n"
|
|
return log_text
|
|
|
|
def _parse_trainer_metrics(self, log_data):
|
|
import re
|
|
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
|
|
|
current_epoch = 0
|
|
current_step = None
|
|
|
|
# Iterate over each line in the log
|
|
for line in log_data.split('\n'):
|
|
line = ansi_escape.sub('', line)
|
|
|
|
# Parse avg losses per epoch
|
|
if "avg_loss:" in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
self.avg_loss.append((current_epoch, value))
|
|
|
|
if "avg_loss_text_ce:" in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
self.avg_loss_text_ce.append((current_epoch, value))
|
|
|
|
if "avg_loss_mel_ce:" in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
self.avg_loss_mel_ce.append((current_epoch, value))
|
|
|
|
# Parse step metrics
|
|
if "loss:" in line and "avg_" not in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
self.loss.append((current_step, value))
|
|
|
|
if "loss_text_ce:" in line and "avg_" not in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
self.loss_text_ce.append((current_step, value))
|
|
|
|
if "loss_mel_ce:" in line and "avg_" not in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
self.loss_mel_ce.append((current_step, value))
|
|
|
|
# Parse learning rate
|
|
if "current_lr:" in line:
|
|
value = float(line.split(":")[1].strip())
|
|
self.learning_rates.append((current_step, value))
|
|
|
|
# Parse gradient norms
|
|
if "grad_norm:" in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
if not hasattr(self, 'grad_norms'):
|
|
self.grad_norms = []
|
|
self.grad_norms.append((current_step, value))
|
|
|
|
# Parse step and loader times
|
|
if "step_time:" in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
if not hasattr(self, 'step_times'):
|
|
self.step_times = []
|
|
self.step_times.append((current_step, value))
|
|
|
|
if "loader_time:" in line:
|
|
value = float(line.split(":")[1].split()[0].strip())
|
|
if not hasattr(self, 'loader_times'):
|
|
self.loader_times = []
|
|
self.loader_times.append((current_step, value))
|
|
|
|
# Parse training and validation loss per epoch
|
|
if "training_loss:" in line:
|
|
value = float(line.split(":")[1].strip())
|
|
if not hasattr(self, 'training_losses'):
|
|
self.training_losses = []
|
|
self.training_losses.append((current_epoch, value))
|
|
|
|
if "validation_loss:" in line:
|
|
value = float(line.split(":")[1].strip())
|
|
if not hasattr(self, 'validation_losses'):
|
|
self.validation_losses = []
|
|
self.validation_losses.append((current_epoch, value))
|
|
|
|
# Update epoch and step trackers
|
|
if "EPOCH:" in line:
|
|
current_epoch = int(line.split('/')[0].split(':')[-1].strip())
|
|
|
|
if "GLOBAL_STEP:" in line:
|
|
current_step = int(line.split("GLOBAL_STEP:")[1].split()[0])
|
|
|
|
def plot_metrics(self):
|
|
# Get the directory of the current script
|
|
current_script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
# Define the full path for the output PNG
|
|
output_path = os.path.join(current_script_dir, "training_metrics.png")
|
|
# Ensure log data is parsed before plotting
|
|
if self.current_model_path:
|
|
log_file = os.path.join(self.current_model_path, "trainer_0_log.txt")
|
|
if os.path.exists(log_file):
|
|
try:
|
|
with open(log_file, 'r') as f:
|
|
log_data = f.read()
|
|
# print(f"[DEBUG] Successfully read {len(log_data)} characters from log file.")
|
|
# Explicitly call the parsing method to process log data
|
|
self._parse_trainer_metrics(log_data)
|
|
except Exception as e:
|
|
print(f"[DEBUG] Error reading trainer log: {e}")
|
|
return output_path # Return early if there's an issue reading the file
|
|
else:
|
|
print("[DEBUG] Log file not found. No data to plot.")
|
|
return output_path # Return early if log file is missing
|
|
plt.rcParams.update({'font.size': 12})
|
|
plt.style.use('dark_background')
|
|
|
|
fig, axes = plt.subplots(3, 2, figsize=(24, 18))
|
|
(ax1, ax2), (ax3, ax4), (ax5, ax6) = axes
|
|
|
|
# Common styling for all axes
|
|
for ax in axes.flat:
|
|
ax.grid(True, alpha=0.3)
|
|
ax.set_facecolor('black')
|
|
for spine in ax.spines.values():
|
|
spine.set_color('white')
|
|
|
|
# 1. Epoch Metrics (looks good but add y-axis adjustment)
|
|
if self.avg_loss or self.avg_loss_text_ce or self.avg_loss_mel_ce:
|
|
ax1.set_title('Epoch Metrics (Avg Losses)')
|
|
ax1.set_xlabel('Epoch')
|
|
ax1.set_ylabel('Loss Value')
|
|
if self.avg_loss:
|
|
epochs, values = zip(*self.avg_loss)
|
|
ax1.plot(epochs, values, color='green', linestyle='-', marker='o', markersize=8, label='Avg Loss')
|
|
if self.avg_loss_text_ce:
|
|
epochs, values = zip(*self.avg_loss_text_ce)
|
|
ax1.plot(epochs, values, color='red', linestyle='-', marker='o', markersize=8, label='Avg Loss Text CE')
|
|
if self.avg_loss_mel_ce:
|
|
epochs, values = zip(*self.avg_loss_mel_ce)
|
|
ax1.plot(epochs, values, color='blue', linestyle='-', marker='o', markersize=8, label='Avg Loss MEL CE')
|
|
ax1.set_xlim(left=0)
|
|
ax1.legend()
|
|
|
|
# 2. Step-wise Loss Metrics (fix x-axis)
|
|
if self.loss or self.loss_text_ce or self.loss_mel_ce:
|
|
ax2.set_title('Step-wise Loss Metrics')
|
|
ax2.set_xlabel('Step')
|
|
ax2.set_ylabel('Loss Value')
|
|
if self.loss:
|
|
steps, values = zip(*self.loss)
|
|
ax2.plot(steps, values, color='green', linestyle='-', marker='o', markersize=8, label='Loss')
|
|
if self.loss_text_ce:
|
|
steps, values = zip(*self.loss_text_ce)
|
|
ax2.plot(steps, values, color='red', linestyle='-', marker='o', markersize=8, label='Loss Text CE')
|
|
if self.loss_mel_ce:
|
|
steps, values = zip(*self.loss_mel_ce)
|
|
ax2.plot(steps, values, color='blue', linestyle='-', marker='o', markersize=8, label='Loss MEL CE')
|
|
ax2.set_xlim(left=0)
|
|
ax2.legend()
|
|
|
|
# 3. Learning Rate Schedule (fix x-axis and scale)
|
|
if self.learning_rates:
|
|
ax3.set_title('Learning Rate Schedule')
|
|
ax3.set_xlabel('Step')
|
|
ax3.set_ylabel('Learning Rate Value')
|
|
steps, values = zip(*self.learning_rates)
|
|
ax3.plot(steps, values, color='yellow', linestyle='-', marker='o', markersize=8, label='Learning Rate')
|
|
ax3.set_xlim(left=0)
|
|
ax3.legend()
|
|
|
|
# 4. Gradient Norm
|
|
if hasattr(self, 'grad_norms') and self.grad_norms:
|
|
ax4.set_title('Gradient Norm over Steps')
|
|
ax4.set_xlabel('Step')
|
|
ax4.set_ylabel('Gradient Norm')
|
|
ax4.plot(*zip(*self.grad_norms), color='purple', linestyle='-', marker='o', markersize=8, label='Gradient Norm')
|
|
ax4.legend()
|
|
|
|
# 5. Step and Loader Times
|
|
if hasattr(self, 'step_times') and self.step_times:
|
|
ax5.set_title('Step and Loader Times')
|
|
ax5.set_xlabel('Step')
|
|
ax5.set_ylabel('Time (seconds)')
|
|
ax5.plot(*zip(*self.step_times), color='orange', linestyle='-', marker='o', markersize=8, label='Step Time')
|
|
if hasattr(self, 'loader_times') and self.loader_times:
|
|
ax5.plot(*zip(*self.loader_times), color='cyan', linestyle='-', marker='o', markersize=8, label='Loader Time')
|
|
ax5.legend()
|
|
|
|
# 6. Training vs Validation Loss
|
|
if hasattr(self, 'training_losses') and self.training_losses:
|
|
ax6.set_title('Training vs Validation Loss')
|
|
ax6.set_xlabel('Epoch')
|
|
ax6.set_ylabel('Loss Value')
|
|
ax6.plot(*zip(*self.training_losses), color='green', linestyle='-', marker='o', markersize=8, label='Training Loss')
|
|
if hasattr(self, 'validation_losses') and self.validation_losses:
|
|
ax6.plot(*zip(*self.validation_losses), color='red', linestyle='-', marker='o', markersize=8, label='Validation Loss')
|
|
ax6.legend()
|
|
|
|
plt.tight_layout()
|
|
for ax in axes.flat:
|
|
ax.set_xlim(left=0) # Ensure no negative x values
|
|
ax.grid(True, alpha=0.3) # Consistent grid
|
|
|
|
# Save the file
|
|
plt.savefig(output_path, format='png', bbox_inches='tight', pad_inches=0.1)
|
|
plt.close()
|
|
|
|
# Return the full path to the calling script
|
|
return output_path |