code: def test_topk(self, k):
      ctx = mx.gpu(0)
      net = get_model("resnet18_v2")
      net.initialize(mx.init.Xavier(), ctx=ctx)
      net.summary(nd.ones((1, 3, 224, 224), ctx=ctx))
  
      # hyper-params
      batch_size = 32
      optimizer_params = {"momentum": 0, "wd": 0, "learning_rate": 0.01}
  
      compression_params = {
          "compressor": "topk",
          "k": k,
      }
  
      trainer = bps.DistributedTrainer(
          net.collect_params(),
          "sgd",
          optimizer_params,
          compression_params=compression_params,
      )
  
      loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()
  
      train_data = fake_data(batch_size=batch_size)
  
      params = {}
  
      for i, param in enumerate(trainer._params):
          if param.grad_req != "null":
              params[i] = param._data[0].asnumpy()
  
      for it, batch in tqdm(enumerate(train_data)):
          data = batch[0].as_in_context(ctx)
          label = batch[1].as_in_context(ctx)
  
          with autograd.record():
              output = net(data)
              loss = loss_fn(output, label)
  
          loss.backward()
  
          gs = {}
          xs = {}
  
          for i, param in enumerate(trainer._params):
              if param.grad_req != "null":
                  gs[i] = param._grad[0].asnumpy()
                  xs[i] = param._data[0].asnumpy()
  
          trainer.step(batch_size)
  
          for i, param in enumerate(trainer._params):
              if param.grad_req != "null":
                  g = gs[i] / (batch_size * bps.size())
                  c = topk(g, k)
  
                  cs = topk(c, k)
                  c = cs
  
                  params[i] -= optimizer_params["learning_rate"] * c
  
      cnt = 0
      tot = 0
      for i, param in enumerate(trainer._params):
          if param.grad_req != "null":
              x = param._data[0].asnumpy()
              tot += len(x.flatten())
              if not np.allclose(
                  params[i], x, atol=np.finfo(np.float32).eps
              ):
                  diff = np.abs(x.flatten() - params[i].flatten())
                  idx = np.where(diff > np.finfo(np.float32).eps)
                  cnt += len(idx[0])
  
      assert cnt == 0, "false/tot=%d/%d=%f" % (cnt, tot, cnt / tot)
  
file_path: byteps/tests/test_topk.py
imports: import itertools
  import random
  import unittest
  import byteps.mxnet
  import mxnet
  import mxnet.ndarray
  import numpy
  from gluoncv.model_zoo import get_model
  from mxnet import autograd, gluon
  from parameterized import parameterized
  from tqdm import tqdm
  from meta_test import MetaTest
  from utils import fake_data
label_desc: Use image node_id 1 for calling the TopkTestCase obj's underlying member method code with example usage: obj.test_topk(k) without return types