code: def main(args):
      np.random.seed(args.seed)
      mx.random.seed(args.seed)
  
      best_epoch = -1
      best_dev_acc = 0
  
      cuda = args.gpu >= 0
      if cuda:
          if args.gpu in mx.test_utils.list_gpus():
              ctx = mx.gpu(args.gpu)
          else:
              print(
                  "Requested GPU id {} was not found. Defaulting to CPU implementation".format(
                      args.gpu
                  )
              )
              ctx = mx.cpu()
      else:
          ctx = mx.cpu()
  
      if args.use_glove:
          prepare_glove()
  
      trainset = data.SSTDataset()
      train_loader = gluon.data.DataLoader(
          dataset=trainset,
          batch_size=args.batch_size,
          batchify_fn=batcher(ctx),
          shuffle=True,
          num_workers=0,
      )
      devset = data.SSTDataset(mode="dev")
      dev_loader = gluon.data.DataLoader(
          dataset=devset,
          batch_size=100,
          batchify_fn=batcher(ctx),
          shuffle=True,
          num_workers=0,
      )
  
      testset = data.SSTDataset(mode="test")
      test_loader = gluon.data.DataLoader(
          dataset=testset,
          batch_size=100,
          batchify_fn=batcher(ctx),
          shuffle=False,
          num_workers=0,
      )
  
      model = TreeLSTM(
          trainset.vocab_size,
          args.x_size,
          args.h_size,
          trainset.num_classes,
          args.dropout,
          cell_type="childsum" if args.child_sum else "nary",
          pretrained_emb=trainset.pretrained_emb,
          ctx=ctx,
      )
      print(model)
      params_ex_emb = [
          x
          for x in model.collect_params().values()
          if x.grad_req != "null" and x.shape[0] != trainset.vocab_size
      ]
      params_emb = list(model.embedding.collect_params().values())
      for p in params_emb:
          p.lr_mult = 0.1
  
      model.initialize(mx.init.Xavier(magnitude=1), ctx=ctx)
      model.hybridize()
      trainer = gluon.Trainer(
          model.collect_params("^(?!embedding).*$"),
          "adagrad",
          {"learning_rate": args.lr, "wd": args.weight_decay},
      )
      trainer_emb = gluon.Trainer(
          model.collect_params("^embedding.*$"),
          "adagrad",
          {"learning_rate": args.lr},
      )
  
      dur = []
      L = gluon.loss.SoftmaxCrossEntropyLoss(axis=1)
      for epoch in range(args.epochs):
          t_epoch = time.time()
          for step, batch in enumerate(train_loader):
              g = batch.graph
              n = g.number_of_nodes()
  
              # TODO begin_states function?
              h = mx.nd.zeros((n, args.h_size), ctx=ctx)
              c = mx.nd.zeros((n, args.h_size), ctx=ctx)
              if step >= 3:
                  t0 = time.time()  # tik
              with mx.autograd.record():
                  pred = model(batch, h, c)
                  loss = L(pred, batch.label)
  
              loss.backward()
              trainer.step(args.batch_size)
              trainer_emb.step(args.batch_size)
  
              if step >= 3:
                  dur.append(time.time() - t0)  # tok
  
              if step > 0 and step % args.log_every == 0:
                  pred = pred.argmax(axis=1).astype(batch.label.dtype)
                  acc = (batch.label == pred).sum()
                  root_ids = [
                      i
                      for i in range(batch.graph.number_of_nodes())
                      if batch.graph.out_degrees(i) == 0
                  ]
                  root_acc = np.sum(
                      batch.label.asnumpy()[root_ids]
                      == pred.asnumpy()[root_ids]
                  )
  
                  print(
                      "Epoch {:05d} | Step {:05d} | Loss {:.4f} | Acc {:.4f} | Root Acc {:.4f} | Time(s) {:.4f}".format(
                          epoch,
                          step,
                          loss.sum().asscalar(),
                          1.0 * acc.asscalar() / len(batch.label),
                          1.0 * root_acc / len(root_ids),
                          np.mean(dur),
                      )
                  )
          print(
              "Epoch {:05d} training time {:.4f}s".format(
                  epoch, time.time() - t_epoch
              )
          )
  
          # eval on dev set
          accs = []
          root_accs = []
          for step, batch in enumerate(dev_loader):
              g = batch.graph
              n = g.number_of_nodes()
              h = mx.nd.zeros((n, args.h_size), ctx=ctx)
              c = mx.nd.zeros((n, args.h_size), ctx=ctx)
              pred = (
                  model(batch, h, c).argmax(1).astype(batch.label.dtype)
              )
  
              acc = (batch.label == pred).sum().asscalar()
              accs.append([acc, len(batch.label)])
              root_ids = [
                  i
                  for i in range(batch.graph.number_of_nodes())
                  if batch.graph.out_degrees(i) == 0
              ]
              root_acc = np.sum(
                  batch.label.asnumpy()[root_ids]
                  == pred.asnumpy()[root_ids]
              )
              root_accs.append([root_acc, len(root_ids)])
  
          dev_acc = (
              1.0
              * np.sum([x[0] for x in accs])
              / np.sum([x[1] for x in accs])
          )
          dev_root_acc = (
              1.0
              * np.sum([x[0] for x in root_accs])
              / np.sum([x[1] for x in root_accs])
          )
          print(
              "Epoch {:05d} | Dev Acc {:.4f} | Root Acc {:.4f}".format(
                  epoch, dev_acc, dev_root_acc
              )
          )
  
          if dev_root_acc > best_dev_acc:
              best_dev_acc = dev_root_acc
              best_epoch = epoch
              model.save_parameters("best_{}.params".format(args.seed))
          else:
              if best_epoch <= epoch - 10:
                  break
  
          # lr decay
          trainer.set_learning_rate(
              max(1e-5, trainer.learning_rate * 0.99)
          )
          print(trainer.learning_rate)
          trainer_emb.set_learning_rate(
              max(1e-5, trainer_emb.learning_rate * 0.99)
          )
          print(trainer_emb.learning_rate)
  
      # test
      model.load_parameters("best_{}.params".format(args.seed))
      accs = []
      root_accs = []
      for step, batch in enumerate(test_loader):
          g = batch.graph
          n = g.number_of_nodes()
          h = mx.nd.zeros((n, args.h_size), ctx=ctx)
          c = mx.nd.zeros((n, args.h_size), ctx=ctx)
          pred = (
              model(batch, h, c)
              .argmax(axis=1)
              .astype(batch.label.dtype)
          )
  
          acc = (batch.label == pred).sum().asscalar()
          accs.append([acc, len(batch.label)])
          root_ids = [
              i
              for i in range(batch.graph.number_of_nodes())
              if batch.graph.out_degrees(i) == 0
          ]
          root_acc = np.sum(
              batch.label.asnumpy()[root_ids]
              == pred.asnumpy()[root_ids]
          )
          root_accs.append([root_acc, len(root_ids)])
  
      test_acc = (
          1.0
          * np.sum([x[0] for x in accs])
          / np.sum([x[1] for x in accs])
      )
      test_root_acc = (
          1.0
          * np.sum([x[0] for x in root_accs])
          / np.sum([x[1] for x in root_accs])
      )
      print(
          "------------------------------------------------------------------------------------"
      )
      print(
          "Epoch {:05d} | Test Acc {:.4f} | Root Acc {:.4f}".format(
              best_epoch, test_acc, test_root_acc
          )
      )
  
file_path: dgl/examples/mxnet/tree_lstm/train.py
imports: import argparse
  import collections
  import os
  import time
  import warnings
  import zipfile
  import dgl
  import dgl.data
  import mxnet
  import numpy
  from mxnet import gluon
  from tree_lstm import TreeLSTM
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types