Lint clean & documentation
Cleaned finetune lint Moved documentation out Moved metrics logger
This commit is contained in:
erew123 authored and GitHub committed
1 parent
b538564da6
commit
adc9bb8cc6
5 files changed
+7283
-5934
No files matched your search
+3648
-3690
File diff suppressed because it is too large.
Load diff
@@ -1,35 +1,35 @@
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
|
||||
class WarmUpScheduler(LRScheduler):
|
||||
def __init__(self, optimizer, total_epochs, warmup_lr, target_lr, after_scheduler=None):
|
||||
self.total_epochs = total_epochs
|
||||
self.warmup_epochs = int(total_epochs * 0.05)
|
||||
if self.warmup_epochs <= 0:
|
||||
self.warmup_epochs = 1
|
||||
self.warmup_lr = warmup_lr
|
||||
self.target_lr = target_lr
|
||||
self.after_scheduler = after_scheduler
|
||||
print(f"[WARMUP] {self.warmup_epochs} epochs")
|
||||
super(WarmUpScheduler, self).__init__(optimizer)
|
||||
|
||||
def get_lr(self):
|
||||
if self.last_epoch < self.warmup_epochs:
|
||||
progress = self.last_epoch / self.warmup_epochs
|
||||
return [self.warmup_lr + progress * (self.target_lr - self.warmup_lr) for _ in self.base_lrs]
|
||||
if self.after_scheduler:
|
||||
return self.after_scheduler.get_last_lr()
|
||||
return [base_lr for base_lr in self.base_lrs]
|
||||
|
||||
def step(self, epoch=None):
|
||||
if epoch is not None and epoch < self.warmup_epochs:
|
||||
self.last_epoch = epoch
|
||||
elif self.after_scheduler and epoch is not None:
|
||||
if epoch < self.warmup_epochs:
|
||||
self.last_epoch = epoch
|
||||
else:
|
||||
self.after_scheduler.step(epoch - self.warmup_epochs)
|
||||
self.last_epoch = epoch
|
||||
else:
|
||||
super(WarmUpScheduler, self).step(epoch)
|
||||
if self.after_scheduler and self.last_epoch >= self.warmup_epochs:
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
|
||||
class WarmUpScheduler(LRScheduler):
|
||||
def __init__(self, optimizer, total_epochs, warmup_lr, target_lr, after_scheduler=None):
|
||||
self.total_epochs = total_epochs
|
||||
self.warmup_epochs = int(total_epochs * 0.05)
|
||||
if self.warmup_epochs <= 0:
|
||||
self.warmup_epochs = 1
|
||||
self.warmup_lr = warmup_lr
|
||||
self.target_lr = target_lr
|
||||
self.after_scheduler = after_scheduler
|
||||
print(f"[WARMUP] {self.warmup_epochs} epochs")
|
||||
super(WarmUpScheduler, self).__init__(optimizer)
|
||||
|
||||
def get_lr(self):
|
||||
if self.last_epoch < self.warmup_epochs:
|
||||
progress = self.last_epoch / self.warmup_epochs
|
||||
return [self.warmup_lr + progress * (self.target_lr - self.warmup_lr) for _ in self.base_lrs]
|
||||
if self.after_scheduler:
|
||||
return self.after_scheduler.get_last_lr()
|
||||
return [base_lr for base_lr in self.base_lrs]
|
||||
|
||||
def step(self, epoch=None):
|
||||
if epoch is not None and epoch < self.warmup_epochs:
|
||||
self.last_epoch = epoch
|
||||
elif self.after_scheduler and epoch is not None:
|
||||
if epoch < self.warmup_epochs:
|
||||
self.last_epoch = epoch
|
||||
else:
|
||||
self.after_scheduler.step(epoch - self.warmup_epochs)
|
||||
self.last_epoch = epoch
|
||||
else:
|
||||
super(WarmUpScheduler, self).step(epoch)
|
||||
if self.after_scheduler and self.last_epoch >= self.warmup_epochs:
|
||||
self.after_scheduler.step(epoch - self.warmup_epochs if epoch is not None else None)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,351 @@
|
||||
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
|
||||
+2210
-2210
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user