MXNet HybridBlock模型并行性

时间:2018-05-06 14:55:37

标签: deep-learning mxnet

我是MXNet的新手,并且一直在试验模型并行性。我找到了这篇不错的帖子:simple example of mxnet model parallelism

并修改代码以使用HybridBlock,如下所示:

import numpy as np
import mxnet as mx
from mxnet import nd, autograd, gluon
from mxnet.gluon import HybridBlock

num_inputs = 2
num_outputs = 1
num_examples = 10000


def real_fn(x):
    return 2 * x[:, 0] - 3.4 * x[:, 1] + 4.2


x = np.random.normal(0, 1, (num_examples, num_inputs))
noise = 0.001 * np.random.normal(0, 1, num_examples)
y = real_fn(x) + noise
y = y.reshape(-1, 1)

hidden_layers = 2
num_gpus = hidden_layers + 1
ctxList = [mx.gpu(i) for i in range(num_gpus)]


class MyDenseBlock(HybridBlock):
    def __init__(self, layer_number, size_input, size_output, **kwargs):
        super(MyDenseBlock, self).__init__(**kwargs)

        self.layer_number = layer_number
        self.size_input = size_input
        self.size_output = size_output

        with self.name_scope():
            # add parameters to the Block's ParameterDict.
            self.weight = self.params.get(
                'weight',
                init=mx.init.Xavier(magnitude=2.24),
                shape=(size_input, size_output),
                grad_req='write')

            self.bias = self.params.get(
                'bias',
                init=mx.init.Constant(0.5),
                shape=(size_output,),
                grad_req='write')

    def hybrid_forward(self, F, x, weight, bias):
        x = x.as_in_context(ctxList[self.layer_number])
        with x.context:
            linear = F.broadcast_add(F.dot(x, weight), bias)
            return linear


net = gluon.nn.HybridSequential()
with net.name_scope():
    net.add(MyDenseBlock(0, size_input=2, size_output=2))

    for i in range(hidden_layers - 1):
        net.add(MyDenseBlock(i + 1, size_input=2, size_output=2))

    net.add(MyDenseBlock(i + 2, size_input=2, size_output=1))


print("\ninitializing:")
params = net.collect_params()

for i, param in enumerate(params):
    if 'mydenseblock0' in param:
        params[param].initialize(ctx=ctxList[0])

    elif 'mydenseblock1' in param:
        params[param].initialize(ctx=ctxList[1])

    elif 'mydenseblock2' in param:
        params[param].initialize(ctx=ctxList[2])

    print("  ", i, param, "  ", params[param].list_data()[0].context)

#net.hybridize()


def square_loss(yhat, y):
    return nd.mean((yhat - y) ** 2)


def custom_trainer(updaters, params, ignore_stale_grad=False):
    for i, param in enumerate(params):
        if params[param].grad_req == 'null':
            continue
        if not ignore_stale_grad:
            for data in params[param].list_data():
                if not data._fresh_grad:
                    print("`%s` on context %s has not been updated" % (params[param].name, str(data.context)))
                    assert False

        for upd, arr, grad in zip(updaters, params[param].list_data(), params[param].list_grad()):
            if not ignore_stale_grad or arr._fresh_grad:
                upd(i, grad, arr)
                arr._fresh_grad = False


batch_size = 100
epochs = 100
iteration = -1

opt = mx.optimizer.create('adam', learning_rate=0.001, rescale_grad=1 / batch_size)
updaters = [mx.optimizer.get_updater(opt)]


results = []
for e in range(epochs):
    train_groups = np.array_split(np.arange(x.shape[0]), x.shape[0] / batch_size)
    for i, idx in enumerate(train_groups):
        iteration += 1
        xtrain, ytrain = x[idx, :], y[idx]

        xtrain = nd.array(xtrain)
        xtrain = xtrain.as_in_context(ctxList[0])

        ytrain = nd.array(ytrain).reshape((-1, 1))
        ytrain = ytrain.as_in_context(ctxList[0])

        with autograd.record():
            yhat = net(xtrain)
            loss = square_loss(yhat, ytrain.as_in_context(ctxList[-1]))

        loss.backward()
        custom_trainer(updaters, net.collect_params())

        if iteration % 10 == 0:
            results.append([iteration, loss.asnumpy().item()])
            print("epoch= {:5,d}, iter= {:6,d},  error= {:6.3E}".format(e, iteration, loss.asnumpy().item()))

但是,我收到以下错误:

RuntimeError: Parameter 'hybridsequential0_mydenseblock1_weight' was 
not initialized on context gpu(0). It was only initialized on [gpu(1)].
terminate called recursively
terminate called after throwing an instance of 'dmlc::Error'

我使用Blocks没有问题(已经在帖子中确认过),只是HybridBlock不起作用。有人可以帮忙吗?在我看来,模型并行性的例子很少。

1 个答案:

答案 0 :(得分:2)

请您告诉我您要使用模型并行性解决的问题是什么?鉴于最近的GPU,大多数模型实际上并不需要模型并行性。尝试思考是否可以使用不同的方法解决问题,如:

  • 使用FP16。
  • 减少模型中参数/图层的数量。
  • 减小输入的大小,例如输入功能的数量。
  • 减少批量大小并进行数据并行化。
  • 将问题分解为可以单独训练的子模型。

如果上述方法都不适用于您并且您确信您确实需要模型并行性,那么以下是Gluon中的模型并行度示例代码:https://github.com/indhub/mxnet_tutorials/blob/master/model_parallelism/ModelParallelism.ipynb

如果您不了解该代码的任何部分,请询问,我可以解释。