code: def train():
      optimizer = optax.chain(
          optax.clip_by_global_norm(1.0),
          optax.adamw(
              FLAGS.learning_rate, weight_decay=FLAGS.weight_decay
          ),
      )
  
      @jax.jit
      def update(params, aux, rng, optim_state, inputs):
          rng, new_rng = jax.random.split(rng)
          (loss, new_aux), grads = loss_vag(params, aux, rng, inputs)
          updates, new_optim_state = optimizer.update(
              grads, optim_state, params
          )
          new_params = optax.apply_updates(updates, params)
          return loss, (new_params, new_aux, new_rng, new_optim_state)
  
      train_data_iter = load_textgrid_wav(
          FLAGS.data_dir,
          FLAGS.max_phoneme_seq_len,
          FLAGS.batch_size,
          FLAGS.max_wave_len,
          "train",
      )
      val_data_iter = load_textgrid_wav(
          FLAGS.data_dir,
          FLAGS.max_phoneme_seq_len,
          FLAGS.batch_size,
          FLAGS.max_wave_len,
          "val",
      )
      melfilter = MelFilter(
          FLAGS.sample_rate,
          FLAGS.n_fft,
          FLAGS.mel_dim,
          FLAGS.fmin,
          FLAGS.fmax,
      )
      batch = next(train_data_iter)
      batch = batch._replace(
          mels=melfilter(batch.wavs.astype(jnp.float32) / (2**15))
      )
      params, aux, rng, optim_state = initial_state(optimizer, batch)
      losses = Deque(maxlen=1000)
      val_losses = Deque(maxlen=100)
  
      last_step = -1
  
      # loading latest checkpoint
      ckpt_fn = FLAGS.ckpt_dir / "acoustic_latest_ckpt.pickle"
      if ckpt_fn.exists():
          print("Resuming from latest checkpoint at", ckpt_fn)
          with open(ckpt_fn, "rb") as f:
              dic = pickle.load(f)
              last_step, params, aux, rng, optim_state = (
                  dic["step"],
                  dic["params"],
                  dic["aux"],
                  dic["rng"],
                  dic["optim_state"],
              )
  
      tr = tqdm(
          range(last_step + 1, FLAGS.num_training_steps + 1),
          desc="training",
          total=FLAGS.num_training_steps + 1,
          initial=last_step + 1,
      )
      for step in tr:
          batch = next(train_data_iter)
          loss, (params, aux, rng, optim_state) = update(
              params, aux, rng, optim_state, batch
          )
          losses.append(loss)
  
          if step % 10 == 0:
              val_batch = next(val_data_iter)
              val_loss, val_aux, predicted_mel, gt_mel = val_loss_fn(
                  params, aux, rng, val_batch
              )
              val_losses.append(val_loss)
              attn = jax.device_get(val_aux["acoustic_model"]["attn"])
              predicted_mel = jax.device_get(predicted_mel[0])
              gt_mel = jax.device_get(gt_mel[0])
  
          if step % 1000 == 0:
              loss = sum(losses).item() / len(losses)
              val_loss = sum(val_losses).item() / len(val_losses)
              tr.write(
                  f"step {step}  train loss {loss:.3f}  val loss {val_loss:.3f}"
              )
  
              # saving predicted mels
              plt.figure(figsize=(10, 10))
              plt.subplot(3, 1, 1)
              plt.imshow(predicted_mel.T, origin="lower", aspect="auto")
              plt.subplot(3, 1, 2)
              plt.imshow(gt_mel.T, origin="lower", aspect="auto")
              plt.subplot(3, 1, 3)
              plt.imshow(attn.T, origin="lower", aspect="auto")
              plt.tight_layout()
              plt.savefig(FLAGS.ckpt_dir / f"mel_{step}.png")
              plt.close()
  
              # saving checkpoint
              with open(ckpt_fn, "wb") as f:
                  pickle.dump(
                      {
                          "step": step,
                          "params": params,
                          "aux": aux,
                          "rng": rng,
                          "optim_state": optim_state,
                      },
                      f,
                  )
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/acoustic_trainer.py
imports: import pickle
  from functools import partial
  from typing import Deque
  import haiku
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import optax
  from tqdm.auto import tqdm
  from vietTTS.nat.config import AcousticInput
  from .config import FLAGS, AcousticInput
  from .data_loader import load_textgrid_wav
  from .dsp import MelFilter
  from .model import AcousticModel
  from .utils import print_flags
label_desc: Use image node_id 5 for calling a global function with example usage: train() and returns: loss,  new_params,  new_aux,  new_rng,  new_optim_state