docstr: Graph Isomorphism Network (GIN) variant introduced in baselines
  for OGB graph property prediction datasets
  
  Parameters
  ----------
  data_info : dict
      The information about the input dataset.
  embed_size : int
      Embedding size.
  num_layers : int
      Number of layers.
  dropout : float
      Dropout rate.
  virtual_node : bool
      Whether to use virtual node.
docstr_tok: ["Graph","Isomorphism","Network","(","GIN",")","variant","introduced","in","baselines","for","OGB","graph","property","prediction","datasets","Parameters","--","--","--","--","--","data_info",":","dict","The","information","about","the","input","dataset",".","embed_size",":","int","Embedding","size",".","num_layers",":","int","Number","of","layers",".","dropout",":","float","Dropout","rate",".","virtual_node",":","bool","Whether","to","use","virtual","node","."]
code: def __init__(
      self,
      data_info: dict,
      embed_size: int = 300,
      num_layers: int = 5,
      dropout: float = 0.5,
      virtual_node: bool = False,
  ):
      """Graph Isomorphism Network (GIN) variant introduced in baselines
      for OGB graph property prediction datasets
  
      Parameters
      ----------
      data_info : dict
          The information about the input dataset.
      embed_size : int
          Embedding size.
      num_layers : int
          Number of layers.
      dropout : float
          Dropout rate.
      virtual_node : bool
          Whether to use virtual node.
      """
      super(OGBGGIN, self).__init__()
      self.data_info = data_info
      self.embed_size = embed_size
      self.num_layers = num_layers
      self.virtual_node = virtual_node
  
      if data_info["name"] in ["ogbg-molhiv", "ogbg-molpcba"]:
          self.node_encoder = AtomEncoder(embed_size)
          self.edge_encoders = nn.ModuleList(
              [BondEncoder(embed_size) for _ in range(num_layers)]
          )
      else:
          # Handle other datasets
          self.node_encoder = nn.Linear(
              data_info["node_feat_size"], embed_size
          )
          self.edge_encoders = nn.ModuleList(
              [
                  nn.Linear(data_info["edge_feat_size"], embed_size)
                  for _ in range(num_layers)
              ]
          )
  
      self.conv_layers = nn.ModuleList(
          [GINEConv(MLP(embed_size)) for _ in range(num_layers)]
      )
  
      self.dropout = nn.Dropout(dropout)
      self.pool = AvgPooling()
      self.pred = nn.Linear(embed_size, data_info["out_size"])
  
      if virtual_node:
          self.virtual_emb = nn.Embedding(1, embed_size)
          nn.init.constant_(self.virtual_emb.weight.data, 0)
          self.mlp_virtual = nn.ModuleList()
          for _ in range(num_layers - 1):
              self.mlp_virtual.append(MLP(embed_size))
          self.virtual_pool = SumPooling()
  
file_path: dgl/dglgo/dglgo/model/graph_encoder/gin_ogbg.py
imports: import dgl
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl.nn import AvgPooling, GINEConv, SumPooling
  from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
label_desc: Use image node_id 1 to create a new OGBGGIN object from inherited base classes: nn with example: obj = OGBGGIN(data_info, embed_size, num_layers, dropout, virtual_node)