class_docstr: The base script for running different 3DGN methods.
class_docstr_tok: ["The","base","script","for","running","different","3DGN","methods","."]
docstr: The run script for training and validation.
  
  Args:
      device (torch.device): Device for computation.
      train_dataset: Training data.
      valid_dataset: Validation data.
      test_dataset: Test data.
      model: Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.
      loss_func (function): The used loss funtion for training.
      evaluation (function): The evaluation function. 
      epochs (int, optinal): Number of total training epochs. (default: :obj:`500`)
      batch_size (int, optinal): Number of samples in each minibatch in training. (default: :obj:`32`)
      vt_batch_size (int, optinal): Number of samples in each minibatch in validation/testing. (default: :obj:`32`)
      lr (float, optinal): Initial learning rate. (default: :obj:`0.0005`)
      lr_decay_factor (float, optinal): Learning rate decay factor. (default: :obj:`0.5`)
      lr_decay_step_size (int, optinal): epochs at which lr_initial <- lr_initial * lr_decay_factor. (default: :obj:`50`)
      weight_decay (float, optinal): weight decay factor at the regularization term. (default: :obj:`0`)
      energy_and_force (bool, optional): If set to :obj:`True`, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: :obj:`False`)    
      p (int, optinal): The forces’ weight for a joint loss of forces and conserved energy during training. (default: :obj:`100`)
      save_dir (str, optinal): The path to save trained models. If set to :obj:`''`, will not save the model. (default: :obj:`''`)
      log_dir (str, optinal): The path to save log files. If set to :obj:`''`, will not save the log files. (default: :obj:`''`)
docstr_tok: ["The","run","script","for","training","and","validation",".","Args",":","device","(","torch.device",")",":","Device","for","computation",".","train_dataset",":","Training","data",".","valid_dataset",":","Validation","data",".","test_dataset",":","Test","data",".","model",":","Which","3DGN","model","to","use",".","Should","be","one","of","the","SchNet",",","DimeNetPP",",","and","SphereNet",".","loss_func","(","function",")",":","The","used","loss","funtion","for","training",".","evaluation","(","function",")",":","The","evaluation","function",".","epochs","(","int",",","optinal",")",":","Number","of","total","training","epochs",".","(","default",":",":","obj",":","`","500","`",")","batch_size","(","int",",","optinal",")",":","Number","of","samples","in","each","minibatch","in","training",".","(","default",":",":","obj",":","`","32","`",")","vt_batch_size","(","int",",","optinal",")",":","Number","of","samples","in","each","minibatch","in","validation\/testing",".","(","default",":",":","obj",":","`","32","`",")","lr","(","float",",","optinal",")",":","Initial","learning","rate",".","(","default",":",":","obj",":","`","0.0005","`",")","lr_decay_factor","(","float",",","optinal",")",":","Learning","rate","decay","factor",".","(","default",":",":","obj",":","`","0.5","`",")","lr_decay_step_size","(","int",",","optinal",")",":","epochs","at","which","lr_initial","<","-","lr_initial","*","lr_decay_factor",".","(","default",":",":","obj",":","`","50","`",")","weight_decay","(","float",",","optinal",")",":","weight","decay","factor","at","the","regularization","term",".","(","default",":",":","obj",":","`","0","`",")","energy_and_force","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","will","predict","energy","and","take","the","minus","derivative","of","the","energy","with","respect","to","the","atomic","positions","as","predicted","forces",".","(","default",":",":","obj",":","`","False","`",")","p","(","int",",","optinal",")",":","The","forces","\u2019","weight","for","a","joint","loss","of","forces","and","conserved","energy","during","training",".","(","default",":",":","obj",":","`","100","`",")","save_dir","(","str",",","optinal",")",":","The","path","to","save","trained","models",".","If","set","to",":","obj",":","`","``","`",",","will","not","save","the","model",".","(","default",":",":","obj",":","`","``","`",")","log_dir","(","str",",","optinal",")",":","The","path","to","save","log","files",".","If","set","to",":","obj",":","`","``","`",",","will","not","save","the","log","files",".","(","default",":",":","obj",":","`","``","`",")"]
code: def run(
      self,
      device,
      train_dataset,
      valid_dataset,
      test_dataset,
      model,
      loss_func,
      evaluation,
      epochs=500,
      batch_size=32,
      vt_batch_size=32,
      lr=0.0005,
      lr_decay_factor=0.5,
      lr_decay_step_size=50,
      weight_decay=0,
      energy_and_force=False,
      p=100,
      save_dir="",
      log_dir="",
  ):
      r"""
      The run script for training and validation.
  
      Args:
          device (torch.device): Device for computation.
          train_dataset: Training data.
          valid_dataset: Validation data.
          test_dataset: Test data.
          model: Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.
          loss_func (function): The used loss funtion for training.
          evaluation (function): The evaluation function.
          epochs (int, optinal): Number of total training epochs. (default: :obj:`500`)
          batch_size (int, optinal): Number of samples in each minibatch in training. (default: :obj:`32`)
          vt_batch_size (int, optinal): Number of samples in each minibatch in validation/testing. (default: :obj:`32`)
          lr (float, optinal): Initial learning rate. (default: :obj:`0.0005`)
          lr_decay_factor (float, optinal): Learning rate decay factor. (default: :obj:`0.5`)
          lr_decay_step_size (int, optinal): epochs at which lr_initial <- lr_initial * lr_decay_factor. (default: :obj:`50`)
          weight_decay (float, optinal): weight decay factor at the regularization term. (default: :obj:`0`)
          energy_and_force (bool, optional): If set to :obj:`True`, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: :obj:`False`)
          p (int, optinal): The forces’ weight for a joint loss of forces and conserved energy during training. (default: :obj:`100`)
          save_dir (str, optinal): The path to save trained models. If set to :obj:`''`, will not save the model. (default: :obj:`''`)
          log_dir (str, optinal): The path to save log files. If set to :obj:`''`, will not save the log files. (default: :obj:`''`)
  
      """
  
      model = model.to(device)
      num_params = sum(p.numel() for p in model.parameters())
      print(f"#Params: {num_params}")
      optimizer = Adam(
          model.parameters(), lr=lr, weight_decay=weight_decay
      )
      scheduler = StepLR(
          optimizer, step_size=lr_decay_step_size, gamma=lr_decay_factor
      )
  
      train_loader = DataLoader(train_dataset, batch_size, shuffle=True)
      valid_loader = DataLoader(
          valid_dataset, vt_batch_size, shuffle=False
      )
      test_loader = DataLoader(
          test_dataset, vt_batch_size, shuffle=False
      )
      best_valid = float("inf")
      best_test = float("inf")
  
      if save_dir != "":
          if not os.path.exists(save_dir):
              os.makedirs(save_dir)
      if log_dir != "":
          if not os.path.exists(log_dir):
              os.makedirs(log_dir)
          writer = SummaryWriter(log_dir=log_dir)
  
      for epoch in range(1, epochs + 1):
          print("\n=====Epoch {}".format(epoch), flush=True)
  
          print("\nTraining...", flush=True)
          train_mae = self.train(
              model,
              optimizer,
              train_loader,
              energy_and_force,
              p,
              loss_func,
              device,
          )
  
          print("\n\nEvaluating...", flush=True)
          valid_mae = self.val(
              model,
              valid_loader,
              energy_and_force,
              p,
              evaluation,
              device,
          )
  
          print("\n\nTesting...", flush=True)
          test_mae = self.val(
              model,
              test_loader,
              energy_and_force,
              p,
              evaluation,
              device,
          )
  
          print()
          print(
              {
                  "Train": train_mae,
                  "Validation": valid_mae,
                  "Test": test_mae,
              }
          )
  
          if log_dir != "":
              writer.add_scalar("train_mae", train_mae, epoch)
              writer.add_scalar("valid_mae", valid_mae, epoch)
              writer.add_scalar("test_mae", test_mae, epoch)
  
          if valid_mae < best_valid:
              best_valid = valid_mae
              best_test = test_mae
              if save_dir != "":
                  print("Saving checkpoint...")
                  checkpoint = {
                      "epoch": epoch,
                      "model_state_dict": model.state_dict(),
                      "optimizer_state_dict": optimizer.state_dict(),
                      "scheduler_state_dict": scheduler.state_dict(),
                      "best_valid_mae": best_valid,
                      "num_params": num_params,
                  }
                  torch.save(
                      checkpoint,
                      os.path.join(save_dir, "valid_checkpoint.pt"),
                  )
  
          scheduler.step()
  
      print(f"Best validation MAE so far: {best_valid}")
      print(f"Test MAE when got best validation result: {best_test}")
  
      if log_dir != "":
          writer.close()
  
file_path: DIG/dig/threedgraph/method/run.py
imports: import time
  import os
  import torch
  from torch.optim import Adam
  from torch_geometric.data import DataLoader
  import numpy
  from torch.autograd import grad
  from torch.utils.tensorboard import SummaryWriter
  from torch.optim.lr_scheduler import StepLR
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling the run obj's underlying member method code with example usage: obj.run(device, train_dataset, valid_dataset, test_dataset, model, loss_func, evaluation, epochs, batch_size, vt_batch_size, lr, lr_decay_factor, lr_decay_step_size, weight_decay, energy_and_force, p, save_dir, log_dir) without return types