autogluon_test1.py

import numpy as np
import autogluon as ag
import mxnet as mx
from mxnet.gluon import nn, Trainer
from mxnet.util import use_np


def get_mxnet_visible_gpus():
    """Get the number of GPUs that are visible to MXNet.
    Returns
    -------
    ctx_l
        The ctx list
    """
    import mxnet as mx
    gpu_ctx_l = []
    for i in range(10):
        try:
            arr = mx.np.array(1.0, ctx=mx.gpu(i))
            arr.asnumpy()
            gpu_ctx_l.append(mx.gpu(i))
        except Exception:
            continue
    return gpu_ctx_l


@use_np
class Net:
    def train_fn(self, args, reporter):
        gpu_ctx_l = get_mxnet_visible_gpus()
        print('num_gpus:', len(gpu_ctx_l))
        np.random.seed(123)
        mx.random.seed(123)
        net = nn.HybridSequential()
        net.add(nn.Dense(16))
        net.add(nn.Activation('relu'))
        net.add(nn.Dense(4))
        net.hybridize()
        net.initialize(ctx=gpu_ctx_l)
        trainer = Trainer(net.collect_params(), 'adam')
        for i in range(100):
            with mx.autograd.record():
                data = mx.np.random.normal(0, 1, (8, 4), ctx=gpu_ctx_l[0])
                out = net(data)
                loss = mx.np.square(out - data).sum()
                loss.backward()
                reporter(loss=loss.asnumpy().item(), iteration=i)
            trainer.step(1.0)



def run_tuning_jobs(fn, search_space):
    args_decorator = ag.args(**search_space)
    scheduler = ag.scheduler.FIFOScheduler(args_decorator(fn),
                                       resource={'num_cpus': 4, 'num_gpus': 1},
                                       num_trials=20,
                                       reward_attr='loss',
                                       time_attr='iteration')
    scheduler.run()
    scheduler.join_jobs()
    return scheduler


search_space = {
    'num_hidden': ag.space.Int(16, 32),
    'lr': ag.space.Real(1e-3, 1e-2)
}

net = Net()

scheduler = run_tuning_jobs(net.train_fn, search_space)