from __future__ import annotations
import os
import gc
from tqdm import tqdm
import wandb
import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset, SequentialSampler
from torch.optim.lr_scheduler import LinearLR, SequentialLR
from einops import rearrange
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs
from ema_pytorch import EMA
from model import CFM
from model.utils import exists, default
from model.dataset import DynamicBatchSampler, collate_fn
# trainer
class Trainer:
def __init__(
self,
model: CFM,
epochs,
learning_rate,
num_warmup_updates = 20000,
save_per_updates = 1000,
checkpoint_path = None,
batch_size = 32,
batch_size_type: str = "sample",
max_samples = 32,
grad_accumulation_steps = 1,
max_grad_norm = 1.0,
noise_scheduler: str | None = None,
duration_predictor: torch.nn.Module | None = None,
wandb_project = "test_e2-tts",
wandb_run_name = "test_run",
wandb_resume_id: str = None,
last_per_steps = None,
accelerate_kwargs: dict = dict(),
ema_kwargs: dict = dict()
):
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters = True)
self.accelerator = Accelerator(
log_with = "wandb",
kwargs_handlers = [ddp_kwargs],
gradient_accumulation_steps = grad_accumulation_steps,
**accelerate_kwargs
)
if exists(wandb_resume_id):
init_kwargs={"wandb": {"resume": "allow", "name": wandb_run_name, 'id': wandb_resume_id}}
else:
init_kwargs={"wandb": {"resume": "allow", "name": wandb_run_name}}
self.accelerator.init_trackers(
project_name = wandb_project,
init_kwargs=init_kwargs,
config={"epochs": epochs,
"learning_rate": learning_rate,
"num_warmup_updates": num_warmup_updates,
"batch_size": batch_size,
"batch_size_type": batch_size_type,
"max_samples": max_samples,
"grad_accumulation_steps": grad_accumulation_steps,
"max_grad_norm": max_grad_norm,
"gpus": self.accelerator.num_processes,
"noise_scheduler": noise_scheduler}
)
self.model = model
if self.is_main:
self.ema_model = EMA(
model,
include_online_model = False,
**ema_kwargs
)
self.ema_model.to(self.accelerator.device)
self.epochs = epochs
self.num_warmup_updates = num_warmup_updates
self.save_per_updates = save_per_updates
self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
self.checkpoint_path = default(checkpoint_path, 'ckpts/test_e2-tts')
self.batch_size = batch_size
self.batch_size_type = batch_size_type
self.max_samples = max_samples
self.grad_accumulation_steps = grad_accumulation_steps
self.max_grad_norm = max_grad_norm
self.noise_scheduler = noise_scheduler
self.duration_predictor = duration_predictor
self.optimizer = AdamW(model.parameters(), lr=learning_rate)
self.model, self.optimizer = self.accelerator.prepare(
self.model, self.optimizer
)
@property
def is_main(self):
return self.accelerator.is_main_process
def save_checkpoint(self, step, last=False):
self.accelerator.wait_for_everyone()
if self.is_main:
checkpoint = dict(
model_state_dict = self.accelerator.unwrap_model(self.model).state_dict(),
optimizer_state_dict = self.accelerator.unwrap_model(self.optimizer).state_dict(),
ema_model_state_dict = self.ema_model.state_dict(),
scheduler_state_dict = self.scheduler.state_dict(),
step = step
)
if not os.path.exists(self.checkpoint_path):
os.makedirs(self.checkpoint_path)
if last == True:
self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
print(f"Saved last checkpoint at step {step}")
else:
self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
def load_checkpoint(self):
if not exists(self.checkpoint_path) or not os.path.exists(self.checkpoint_path) or not os.listdir(self.checkpoint_path):
return 0
self.accelerator.wait_for_everyone()
if "model_last.pt" in os.listdir(self.checkpoint_path):
latest_checkpoint = "model_last.pt"
else:
latest_checkpoint = sorted([f for f in os.listdir(self.checkpoint_path) if f.endswith('.pt')], key=lambda x: int(''.join(filter(str.isdigit, x))))[-1]
# checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location=self.accelerator.device) # rather use accelerator.load_state ಥ_ಥ
checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location="cpu")
if self.is_main:
self.ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
if 'step' in checkpoint:
self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint['model_state_dict'])
self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint['optimizer_state_dict'])
if self.scheduler:
self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
step = checkpoint['step']
else:
checkpoint['model_state_dict'] = {k.replace("ema_model.", ""): v for k, v in checkpoint['ema_model_state_dict'].items() if k not in ["initted", "step"]}
self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint['model_state_dict'])
step = 0
del checkpoint; gc.collect()
return step
def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None):
if exists(resumable_with_seed):
generator = torch.Generator()
generator.manual_seed(resumable_with_seed)
else:
generator = None
if self.batch_size_type == "sample":
train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True, persistent_workers=True,
batch_size=self.batch_size, shuffle=True, generator=generator)
elif self.batch_size_type == "frame":
self.accelerator.even_batches = False
sampler = SequentialSampler(train_dataset)
batch_sampler = DynamicBatchSampler(sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False)
train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True, persistent_workers=True,
batch_sampler=batch_sampler)
else:
raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")
# accelerator.prepare() dispatches batches to devices;
# which means the length of dataloader calculated before, should consider the number of devices
warmup_steps = self.num_warmup_updates * self.accelerator.num_processes # consider a fixed warmup steps while using accelerate multi-gpu ddp
# otherwise by default with split_batches=False, warmup steps change with num_processes
total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
decay_steps = total_steps - warmup_steps
warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
self.scheduler = SequentialLR(self.optimizer,
schedulers=[warmup_scheduler, decay_scheduler],
milestones=[warmup_steps])
train_dataloader, self.scheduler = self.accelerator.prepare(train_dataloader, self.scheduler) # actual steps = 1 gpu steps / gpus
start_step = self.load_checkpoint()
global_step = start_step
if exists(resumable_with_seed):
orig_epoch_step = len(train_dataloader)
skipped_epoch = int(start_step // orig_epoch_step)
skipped_batch = start_step % orig_epoch_step
skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
else:
skipped_epoch = 0
for epoch in range(skipped_epoch, self.epochs):
self.model.train()
if exists(resumable_with_seed) and epoch == skipped_epoch:
progress_bar = tqdm(skipped_dataloader, desc=f"Epoch {epoch+1}/{self.epochs}", unit="step", disable=not self.accelerator.is_local_main_process,
initial=skipped_batch, total=orig_epoch_step)
else:
progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{self.epochs}", unit="step", disable=not self.accelerator.is_local_main_process)
for batch in progress_bar:
with self.accelerator.accumulate(self.model):
text_inputs = batch['text']
mel_spec = rearrange(batch['mel'], 'b d n -> b n d')
mel_lengths = batch["mel_lengths"]
# TODO. add duration predictor training
if self.duration_predictor is not None and self.accelerator.is_local_main_process:
dur_loss = self.duration_predictor(mel_spec, lens=batch.get('durations'))
self.accelerator.log({"duration loss": dur_loss.item()}, step=global_step)
loss, cond, pred = self.model(mel_spec, text=text_inputs, lens=mel_lengths, noise_scheduler=self.noise_scheduler)
self.accelerator.backward(loss)
if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
self.optimizer.step()
self.scheduler.step()
self.optimizer.zero_grad()
if self.is_main:
self.ema_model.update()
global_step += 1
if self.accelerator.is_local_main_process:
self.accelerator.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step)
progress_bar.set_postfix(step=str(global_step), loss=loss.item())
if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
self.save_checkpoint(global_step)
if global_step % self.last_per_steps == 0:
self.save_checkpoint(global_step, last=True)
self.accelerator.end_training()