ndarray不支持多GPU加载

例如我使用SigmoidBinaryCrossEntropyLoss的时候,对于其第四个参数pos_weight,在使用多GPU训练的时候会出现

mxnet/mshadow/mshadow/./stream_gpu-inl.h:62: Check failed: e == cudaSuccess CUDA: an illegal memory access was encountered

Check failed: type_ != nullptr: The any container is empty requested=N5mxnet

第一个问题是pos_weight多GPU加载位置不当,第二个是GPU并行计算的时候,我得pos_weight只加载在某一个gpu上,另一个gpu并没有加载导致(我这么认为)

目前的解决方案是,权重各自加载到对应gpu上,对应gpu上并行计算时,用各自的权重

ctx = [mx.gpu(), mx.gpu(1))
pos_weights = []
for ctx_i in ctx:    # 同时加载多个w到每个gpu上
    w = nd.array(pos_weight, ctx=ctx_i)
    pos_weights.append(w)

...

        for i, batch in enumerate(train_iter):
            Xs, ys, batch_size = _get_batch(batch, ctx)
            ls = []
            with autograd.record():
                y_hats = [net(X) for X in Xs]
                ls = [loss(y_hat, y, None, pos_weights[ctx.index(y.context)]) for y_hat, y in zip(y_hats, ys)]  # 这里的pos_weights参数就是我说的需要同时加载到各个gpu上的数据
            for l in ls:
                l.backward()

这样就解决了上述错误,但个人认为这或许不太方便