text: import argparse
  import os
  
  import torch
  import torch.nn as nn
  import yaml
  from torch.utils.data import random_split
  from torch.utils.tensorboard import SummaryWriter
  from tqdm.auto import trange
  
  from data import InfiniteDataLoader, SpeakerDataset, infinite_iterator
  from model import AdaINVC
  
  
  def main(
      config_file: str,
      data_dir: str,
      save_dir: str,
      n_steps: int,
      save_steps: int,
      log_steps: int,
      n_spks: int,
      n_uttrs: int,
  ):
      device = "cuda" if torch.cuda.is_available() else "cpu"
      torch.backends.cudnn.benchmark = True
  
      # Load config
      config = yaml.load(open(config_file, "r"), Loader=yaml.FullLoader)
  
      # Prepare data
      data = SpeakerDataset(data_dir, segment=128, n_uttrs=n_uttrs)
  
      # split train/valid sets
      train_set, valid_set = random_split(
          data, [int(len(data) * 0.8), len(data) - int(len(data) * 0.8)]
      )
  
      # construct loader
      train_loader = InfiniteDataLoader(
          train_set, batch_size=n_spks, shuffle=True, num_workers=8
      )
      valid_loader = InfiniteDataLoader(
          valid_set, batch_size=n_spks, shuffle=True, num_workers=8
      )
  
      # construct iterator
      train_iter = infinite_iterator(train_loader)
      valid_iter = infinite_iterator(valid_loader)
  
      # Build model
      model = AdaINVC(config["Model"]).to(device)
      model = torch.jit.script(model)
  
      # Optimizer
      opt = torch.optim.Adam(
          model.parameters(),
          lr=config["Optimizer"]["lr"],
          betas=(config["Optimizer"]["beta1"], config["Optimizer"]["beta2"]),
          amsgrad=config["Optimizer"]["amsgrad"],
          weight_decay=config["Optimizer"]["weight_decay"],
      )
  
      # Tensorboard logger
      writer = SummaryWriter(save_dir)
      criterion = nn.L1Loss()
      pbar = trange(n_steps, ncols=0)
      valid_steps = 32
  
      for step in pbar:
          # get features
          org_mels = next(train_iter)
          org_mels = org_mels.flatten(0, 1)
          org_mels = org_mels.to(device)
  
          # reconstruction
          mu, log_sigma, emb, rec_mels = model(org_mels)
  
          # compute loss
          rec_loss = criterion(rec_mels, org_mels)
          kl_loss = 0.5 * (log_sigma.exp() + mu ** 2 - 1 - log_sigma).mean()
          rec_lambda = config["Lambda"]["rec"]
          kl_lambda = min(
              config["Lambda"]["kl"] * step / config["Lambda"]["kl_annealing"],
              config["Lambda"]["kl"],
          )
          loss = rec_lambda * rec_loss + kl_lambda * kl_loss
  
          # update parameters
          opt.zero_grad()
          loss.backward()
          grad_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=5)
          opt.step()
  
          # save model and optimizer
          if (step + 1) % save_steps == 0:
              model_path = os.path.join(save_dir, f"model-{step + 1}.ckpt")
              model.cpu()
              model.save(model_path)
              model.to(device)
              opt_path = os.path.join(save_dir, f"opt-{step + 1}.ckpt")
              torch.save(opt.state_dict(), opt_path)
  
          if (step + 1) % log_steps == 0:
              # validation
              model.eval()
              valid_loss = 0
              for _ in range(valid_steps):
                  org_mels = next(valid_iter)
                  org_mels = org_mels.flatten(0, 1)
                  org_mels = org_mels.to(device)
                  mu, log_sigma, emb, rec_mels = model(org_mels)
                  loss = criterion(rec_mels, org_mels)
                  valid_loss += loss.item()
              valid_loss /= valid_steps
              model.train()
  
              # record information
              writer.add_scalar("training/rec_loss", rec_loss, step + 1)
              writer.add_scalar("training/kl_loss", kl_loss, step + 1)
              writer.add_scalar("training/grad_norm", grad_norm, step + 1)
              writer.add_scalar("lambda/kl", kl_lambda, step + 1)
              writer.add_scalar("validation/rec_loss", valid_loss, step + 1)
  
          # update tqdm bar
          pbar.set_postfix({"rec_loss": rec_loss.item(), "kl_loss": kl_loss.item()})
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
      parser.add_argument("config_file", type=str)
      parser.add_argument("data_dir", type=str)
      parser.add_argument("save_dir", type=str)
      parser.add_argument("--n_steps", type=int, default=int(1e6))
      parser.add_argument("--save_steps", type=int, default=5000)
      parser.add_argument("--log_steps", type=int, default=250)
      parser.add_argument("--n_spks", type=int, default=32)
      parser.add_argument("--n_uttrs", type=int, default=4)
      main(**vars(parser.parse_args()))
  