class_docstr: The method class for the JTVAE algorithm proposed in the paper `Junction Tree Variational Autoencoder for Molecular Graph Generation <https://arxiv.org/abs/1802.04364>`_. This class provides interfaces for running random generation with the JTVAE algorithm. Please refer to the `example codes <https://github.com/divelab/DIG/tree/dig-stable/examples/ggraph/JTVAE>`_ for usage examples.
  
  Args:
      list_smiles (list): List of smiles in training data.
      training (boolean): If we are training (as opposed to testing).
      build_vocab (boolean): If we need to build the vocabulary (first time training with this dataset).
      device (torch.device, optional): The device where the model is deployed.
docstr: Train the Junction Tree Variational Autoencoder for the random generation task.
  
  Args:
      loader (MolTreeFolder): The MolTreeFolder loader.
      load_epoch (int): The epoch to load from state dictionary.
      lr (float): The learning rate for training.
      anneal_rate (float): The learning rate annealing.
      clip_norm (float): Clips gradient norm of an iterable of parameters.
      num_epochs (int): The number of training epochs.
      beta (float): The KL regularization weight.
      max_beta (float): The maximum KL regularization weight.
      step_beta (float): The KL regularization weight step size.
      anneal_iter (int): How often to step in annealing the learning rate.
      kl_anneal_iter (int): How often to step in annealing the KL regularization weight.
      print_iter (int): How often to print the iteration statistics.
      save_iter (int): How often to save the iteration statistics.
docstr_tok: ["Train","the","Junction","Tree","Variational","Autoencoder","for","the","random","generation","task",".","Args",":","loader","(","MolTreeFolder",")",":","The","MolTreeFolder","loader",".","load_epoch","(","int",")",":","The","epoch","to","load","from","state","dictionary",".","lr","(","float",")",":","The","learning","rate","for","training",".","anneal_rate","(","float",")",":","The","learning","rate","annealing",".","clip_norm","(","float",")",":","Clips","gradient","norm","of","an","iterable","of","parameters",".","num_epochs","(","int",")",":","The","number","of","training","epochs",".","beta","(","float",")",":","The","KL","regularization","weight",".","max_beta","(","float",")",":","The","maximum","KL","regularization","weight",".","step_beta","(","float",")",":","The","KL","regularization","weight","step","size",".","anneal_iter","(","int",")",":","How","often","to","step","in","annealing","the","learning","rate",".","kl_anneal_iter","(","int",")",":","How","often","to","step","in","annealing","the","KL","regularization","weight",".","print_iter","(","int",")",":","How","often","to","print","the","iteration","statistics",".","save_iter","(","int",")",":","How","often","to","save","the","iteration","statistics","."]
code: def train_rand_gen(
      self,
      loader,
      load_epoch,
      lr,
      anneal_rate,
      clip_norm,
      num_epochs,
      beta,
      max_beta,
      step_beta,
      anneal_iter,
      kl_anneal_iter,
      print_iter,
      save_iter,
  ):
      r"""
      Train the Junction Tree Variational Autoencoder for the random generation task.
  
      Args:
          loader (MolTreeFolder): The MolTreeFolder loader.
          load_epoch (int): The epoch to load from state dictionary.
          lr (float): The learning rate for training.
          anneal_rate (float): The learning rate annealing.
          clip_norm (float): Clips gradient norm of an iterable of parameters.
          num_epochs (int): The number of training epochs.
          beta (float): The KL regularization weight.
          max_beta (float): The maximum KL regularization weight.
          step_beta (float): The KL regularization weight step size.
          anneal_iter (int): How often to step in annealing the learning rate.
          kl_anneal_iter (int): How often to step in annealing the KL regularization weight.
          print_iter (int): How often to print the iteration statistics.
          save_iter (int): How often to save the iteration statistics.
  
      """
      vocab = fast_jtnn.Vocab(self.vocab)
  
      for param in self.vae.parameters():
          if param.dim() == 1:
              nn.init.constant_(param, 0)
          else:
              nn.init.xavier_normal_(param)
  
      if 0 > 0:
          self.vae.load_state_dict(
              torch.load(
                  "checkpoints" + "/model.iter-" + str(load_epoch)
              )
          )
  
      print(
          "Model #Params: %dK"
          % (sum([x.nelement() for x in self.vae.parameters()]) / 1000,)
      )
  
      optimizer = optim.Adam(self.vae.parameters(), lr=lr)
      scheduler = lr_scheduler.ExponentialLR(optimizer, anneal_rate)
      scheduler.step()
  
      def param_norm(m):
          return math.sqrt(
              sum([p.norm().item() ** 2 for p in m.parameters()])
          )
  
      def grad_norm(m):
          return math.sqrt(
              sum(
                  [
                      p.grad.norm().item() ** 2
                      for p in m.parameters()
                      if p.grad is not None
                  ]
              )
          )
  
      total_step = load_epoch
      meters = np.zeros(4)
  
      for epoch in range(num_epochs):
          for batch in loader:
              total_step += 1
              try:
                  self.vae.zero_grad()
                  loss, kl_div, wacc, tacc, sacc = self.vae(batch, beta)
                  loss.backward()
                  nn.utils.clip_grad_norm_(
                      self.vae.parameters(), clip_norm
                  )
                  optimizer.step()
              except Exception as e:
                  print(e)
                  continue
  
              meters = meters + np.array(
                  [kl_div, wacc * 100, tacc * 100, sacc * 100]
              )
  
              if total_step % print_iter == 0:
                  meters /= 50
                  print(
                      "[%d] Beta: %.3f, KL: %.2f, Word: %.2f, Topo: %.2f, Assm: %.2f, PNorm: %.2f, GNorm: %.2f"
                      % (
                          total_step,
                          beta,
                          meters[0],
                          meters[1],
                          meters[2],
                          meters[3],
                          param_norm(self.vae),
                          grad_norm(self.vae),
                      )
                  )
                  sys.stdout.flush()
                  meters *= 0
  
              if total_step % save_iter == 0:
                  torch.save(
                      self.vae.state_dict(),
                      "saved" + "/model.iter-" + str(total_step),
                  )
  
              if total_step % anneal_iter == 0:
                  scheduler.step()
                  print("learning rate: %.6f" % scheduler.get_lr()[0])
  
              if (
                  total_step % kl_anneal_iter == 0
                  and total_step >= anneal_iter
              ):
                  beta = min(max_beta, beta + step_beta)
  
imports: import math
  import sys
  from optparse import OptionParser
  import numpy
  import rdkit
  from rdkit import RDLogger, Chem
  from rdkit.Chem import Descriptors
  from tqdm import tqdm
  import torch
  import torch.nn
  import torch.optim
  import torch.optim.lr_scheduler
  from dig.ggraph.method import Generator
  from .None import fast_jtnn
label_desc: Use image node_id 6 for calling the JTVAE obj's underlying member method code with example usage: obj.train_rand_gen(loader, load_epoch, lr, anneal_rate, clip_norm, num_epochs, beta, max_beta, step_beta, anneal_iter, kl_anneal_iter, print_iter, save_iter) and returns: math, math