def setup_mnist_trainer(self, display_log=False, use_chx=False):
batchsize = 100
n_units = 100
comm = self.communicator
model = L.Classifier(MLP(n_units, 10))
model.to_device(get_device(None, use_chx))
optimizer = chainermn.create_multi_node_optimizer(
chainer.optimizers.Adam(), comm)
optimizer.setup(model)
if comm.rank == 0:
train, test = chainer.datasets.get_mnist()
else:
train, test = None, None
train = chainermn.scatter_dataset(train, comm, shuffle=True)
test = chainermn.scatter_dataset(test, comm, shuffle=True)
train_iter = chainer.iterators.SerialIterator(train, batchsize)
test_iter = chainer.iterators.SerialIterator(test,
batchsize,
repeat=False,
shuffle=False)
updater = training.StandardUpdater(train_iter, optimizer)
return updater, optimizer, train_iter, test_iter, model
def check_create_mnbn_model_sequential(self, use_gpu, use_chx):
size = 3
model = chainer.Sequential(
chainer.links.Convolution2D(None, size, 1, 1, 1, nobias=True),
chainer.links.BatchNormalization(size), chainer.functions.relu)
mnbn_model = chainermn.links.create_mnbn_model(model,
self.communicator)
device = get_device(self.communicator.intra_rank if use_gpu else None,
use_chx)
mnbn_model.to_device(device)
with chainer.using_device(mnbn_model.device):
x = mnbn_model.xp.zeros((1, 1, 1, 1))
mnbn_model(x)
def check_create_mnbn_model_chain_list(self, use_gpu, use_chx):
model = BnChainList(3)
mnbn_model = chainermn.links.create_mnbn_model(model,
self.communicator)
self.assertTrue(isinstance(mnbn_model[0], chainer.links.Convolution2D))
self.assertTrue(
isinstance(mnbn_model[1],
chainermn.links.MultiNodeBatchNormalization))
device = get_device(self.communicator.intra_rank if use_gpu else None,
use_chx)
mnbn_model.to_device(device)
with chainer.using_device(mnbn_model.device):
x = mnbn_model.xp.zeros((1, 1, 1, 1))
mnbn_model(x)
def _test(self, comm, model, use_gpu, use_chx):
if use_gpu:
# Use CuPy's Device class to force call cudaSetDevice()
chainer.cuda.get_device_from_id(comm.intra_rank).use()
device = get_device(comm.intra_rank if use_gpu else None, use_chx)
model.to_device(device)
rank = comm.rank
model.bn1.avg_mean.fill(rank * 1)
model.bn2.avg_mean.fill(rank * 2)
model.bn1.avg_var.fill(rank * 3)
model.bn2.avg_var.fill(rank * 4)
allreduce_persistent = \
chainermn.extensions.AllreducePersistent(model, comm)
allreduce_persistent()
avg_rank = (comm.size - 1) / 2.0
chainer.testing.assert_allclose(model.bn1.avg_mean, avg_rank * 1)
chainer.testing.assert_allclose(model.bn2.avg_mean, avg_rank * 2)
chainer.testing.assert_allclose(model.bn1.avg_var, avg_rank * 3)
chainer.testing.assert_allclose(model.bn2.avg_var, avg_rank * 4)