class_docstr: The method class for G-SphereNet algorithm proposed in the paper `An Autoregressive Flow Model for 3D Molecular Geometry Generation from Scratch <https://openreview.net/forum?id=C03Ajc-NS5W>`_. 
  This class provides interfaces for running training and generation with G-SphereNet algorithm. 
  Please refer to the `example codes <https://github.com/divelab/DIG/tree/dig-stable/examples/ggraph3D/G_SphereNet>`_ for usage examples.
docstr: Running training for random generation task.
  
  Args:
      loader: The data loader for loading training samples. It is supposed to use dig.ggraph3D.dataset.QM93DGEN
          as the dataset class, and apply torch.utils.data.DataLoader to it to form the data loader.
      lr (float): The learning rate for training.
      wd (float): The weight decay factor for training.
      max_epochs (int): The maximum number of training epochs.
      model_conf_dict (dict): The python dict for configuring the model hyperparameters.
      save_interval (int): Indicate the frequency to save the model parameters to .pth files,
          *e.g.*, if save_interval=2, the model parameters will be saved for every 2 training epochs.
      save_dir (str): The directory to save the model parameters.
docstr_tok: ["Running","training","for","random","generation","task",".","Args",":","loader",":","The","data","loader","for","loading","training","samples",".","It","is","supposed","to","use","dig.ggraph3D.dataset.QM93DGEN","as","the","dataset","class",",","and","apply","torch.utils.data.DataLoader","to","it","to","form","the","data","loader",".","lr","(","float",")",":","The","learning","rate","for","training",".","wd","(","float",")",":","The","weight","decay","factor","for","training",".","max_epochs","(","int",")",":","The","maximum","number","of","training","epochs",".","model_conf_dict","(","dict",")",":","The","python","dict","for","configuring","the","model","hyperparameters",".","save_interval","(","int",")",":","Indicate","the","frequency","to","save","the","model","parameters","to",".pth","files",",","*","e.g",".","*",",","if","save_interval=2",",","the","model","parameters","will","be","saved","for","every","2","training","epochs",".","save_dir","(","str",")",":","The","directory","to","save","the","model","parameters","."]
code: def train(
      self,
      loader,
      lr,
      wd,
      max_epochs,
      model_conf_dict,
      checkpoint_path,
      save_interval,
      save_dir,
  ):
      r"""
      Running training for random generation task.
  
      Args:
          loader: The data loader for loading training samples. It is supposed to use dig.ggraph3D.dataset.QM93DGEN
              as the dataset class, and apply torch.utils.data.DataLoader to it to form the data loader.
          lr (float): The learning rate for training.
          wd (float): The weight decay factor for training.
          max_epochs (int): The maximum number of training epochs.
          model_conf_dict (dict): The python dict for configuring the model hyperparameters.
          save_interval (int): Indicate the frequency to save the model parameters to .pth files,
              *e.g.*, if save_interval=2, the model parameters will be saved for every 2 training epochs.
          save_dir (str): The directory to save the model parameters.
      """
      self.get_model(model_conf_dict, checkpoint_path)
      self.model.train()
      optimizer = torch.optim.Adam(
          filter(lambda p: p.requires_grad, self.model.parameters()),
          lr=lr,
          weight_decay=wd,
      )
      ce_loss = torch.nn.BCELoss()
      if not os.path.isdir(save_dir):
          os.mkdir(save_dir)
  
      for epoch in range(1, max_epochs + 1):
          total_loss = 0
          for batch, data_batch in enumerate(loader):
              optimizer.zero_grad()
  
              if model_conf_dict["use_gpu"]:
                  for key in data_batch:
                      data_batch[key] = data_batch[key].to("cuda")
  
              (
                  node_out,
                  focus_score,
                  dist_out,
                  angle_out,
                  torsion_out,
              ) = self.model(data_batch)
  
              ll_node = torch.mean(
                  1 / 2 * (node_out[0] ** 2) - node_out[1]
              )
              ll_dist = torch.mean(
                  1 / 2 * (dist_out[0] ** 2) - dist_out[1]
              )
              ll_angle = torch.mean(
                  1 / 2 * (angle_out[0] ** 2) - angle_out[1]
              )
              ll_torsion = torch.mean(
                  1 / 2 * (torsion_out[0] ** 2) - torsion_out[1]
              )
              cannot_focus = data_batch["cannot_focus"]
              focus_ce = ce_loss(focus_score, cannot_focus)
  
              loss = (
                  ll_node + ll_dist + ll_angle + ll_torsion + focus_ce
              )
              loss.backward()
              optimizer.step()
  
              total_loss += loss.to("cpu").item()
              print(
                  "Training iteration {} | loss {}".format(
                      batch, loss.to("cpu").item()
                  )
              )
  
          avg_loss = total_loss / (batch + 1)
          print("Training | Average loss {}".format(avg_loss))
  
          if epoch % save_interval == 0:
              torch.save(
                  self.model.state_dict(),
                  os.path.join(
                      save_dir, "model_ckpt_{}.pth".format(epoch)
                  ),
              )
  
file_path: DIG/dig/ggraph3D/method/G_SphereNet/gspherenet.py
imports: import os
  import torch
  import numpy
  from .model import SphGen
label_desc: Use image node_id 4 for calling the G_SphereNet obj's underlying member method code with example usage: obj.train(loader, lr, wd, max_epochs, model_conf_dict, checkpoint_path, save_interval, save_dir) without return types