code: def main(args):
      # load graph data
      if args.dataset == "aifb":
          dataset = AIFBDataset()
      elif args.dataset == "mutag":
          dataset = MUTAGDataset()
      elif args.dataset == "bgs":
          dataset = BGSDataset()
      elif args.dataset == "am":
          dataset = AMDataset()
      else:
          raise ValueError()
  
      # preprocessing in cpu
      with tf.device("/cpu:0"):
          # Load from hetero-graph
          hg = dataset[0]
  
          num_rels = len(hg.canonical_etypes)
          category = dataset.predict_category
          num_classes = dataset.num_classes
          train_mask = hg.nodes[category].data.pop("train_mask")
          test_mask = hg.nodes[category].data.pop("test_mask")
          train_idx = tf.squeeze(tf.where(train_mask))
          test_idx = tf.squeeze(tf.where(test_mask))
          labels = hg.nodes[category].data.pop("labels")
  
          # split dataset into train, validate, test
          if args.validation:
              val_idx = train_idx[: len(train_idx) // 5]
              train_idx = train_idx[len(train_idx) // 5 :]
          else:
              val_idx = train_idx
  
          # calculate norm for each edge type and store in edge
          for canonical_etype in hg.canonical_etypes:
              u, v, eid = hg.all_edges(
                  form="all", etype=canonical_etype
              )
              _, inverse_index, count = tf.unique_with_counts(v)
              degrees = tf.gather(count, inverse_index)
              norm = tf.ones(eid.shape[0]) / tf.cast(
                  degrees, tf.float32
              )
              norm = tf.expand_dims(norm, 1)
              hg.edges[canonical_etype].data["norm"] = norm
  
          # get target category id
          category_id = len(hg.ntypes)
          for i, ntype in enumerate(hg.ntypes):
              if ntype == category:
                  category_id = i
  
          # edge type and normalization factor
          g = dgl.to_homogeneous(hg, edata=["norm"])
  
      # check cuda
      if args.gpu < 0:
          device = "/cpu:0"
          use_cuda = False
      else:
          device = "/gpu:{}".format(args.gpu)
          g = g.to(device)
          use_cuda = True
      num_nodes = g.number_of_nodes()
      node_ids = tf.range(num_nodes, dtype=tf.int64)
      edge_norm = g.edata["norm"]
      edge_type = tf.cast(g.edata[dgl.ETYPE], tf.int64)
  
      # find out the target node ids in g
      node_tids = g.ndata[dgl.NTYPE]
      loc = node_tids == category_id
      target_idx = tf.squeeze(tf.where(loc))
  
      # since the nodes are featureless, the input feature is then the node id.
      feats = tf.range(num_nodes, dtype=tf.int64)
  
      with tf.device(device):
          # create model
          model = EntityClassify(
              num_nodes,
              args.n_hidden,
              num_classes,
              num_rels,
              num_bases=args.n_bases,
              num_hidden_layers=args.n_layers - 2,
              dropout=args.dropout,
              use_self_loop=args.use_self_loop,
              use_cuda=use_cuda,
          )
  
          # optimizer
          optimizer = tf.keras.optimizers.Adam(learning_rate=args.lr)
          # training loop
          print("start training...")
          forward_time = []
          backward_time = []
          loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy(
              from_logits=False
          )
          for epoch in range(args.n_epochs):
              t0 = time.time()
              with tf.GradientTape() as tape:
                  logits = model(g, feats, edge_type, edge_norm)
                  logits = tf.gather(logits, target_idx)
                  loss = loss_fcn(
                      tf.gather(labels, train_idx),
                      tf.gather(logits, train_idx),
                  )
                  # Manually Weight Decay
                  # We found Tensorflow has a different implementation on weight decay
                  # of Adam(W) optimizer with PyTorch. And this results in worse results.
                  # Manually adding weights to the loss to do weight decay solves this problem.
                  for weight in model.trainable_weights:
                      loss = loss + args.l2norm * tf.nn.l2_loss(weight)
                  t1 = time.time()
                  grads = tape.gradient(loss, model.trainable_weights)
                  optimizer.apply_gradients(
                      zip(grads, model.trainable_weights)
                  )
                  t2 = time.time()
  
              forward_time.append(t1 - t0)
              backward_time.append(t2 - t1)
              print(
                  "Epoch {:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}".format(
                      epoch, forward_time[-1], backward_time[-1]
                  )
              )
              train_acc = acc(logits, labels, train_idx)
              val_loss = loss_fcn(
                  tf.gather(labels, val_idx), tf.gather(logits, val_idx)
              )
              val_acc = acc(logits, labels, val_idx)
              print(
                  "Train Accuracy: {:.4f} | Train Loss: {:.4f} | Validation Accuracy: {:.4f} | Validation loss: {:.4f}".format(
                      train_acc,
                      loss.numpy().item(),
                      val_acc,
                      val_loss.numpy().item(),
                  )
              )
          print()
  
          logits = model(g, feats, edge_type, edge_norm)
          logits = tf.gather(logits, target_idx)
          test_loss = loss_fcn(
              tf.gather(labels, test_idx), tf.gather(logits, test_idx)
          )
          test_acc = acc(logits, labels, test_idx)
          print(
              "Test Accuracy: {:.4f} | Test loss: {:.4f}".format(
                  test_acc, test_loss.numpy().item()
              )
          )
          print()
  
          print(
              "Mean forward time: {:4f}".format(
                  np.mean(forward_time[len(forward_time) // 4 :])
              )
          )
          print(
              "Mean backward time: {:4f}".format(
                  np.mean(backward_time[len(backward_time) // 4 :])
              )
          )
  
file_path: dgl/examples/tensorflow/rgcn/entity_classify.py
imports: import argparse
  import time
  from functools import partial
  import dgl
  import numpy
  import tensorflow
  from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
  from dgl.nn.tensorflow import RelGraphConv
  from model import BaseRGCN
  from tensorflow.keras import layers
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types