Skip to content

Instantly share code, notes, and snippets.

@keisukefukuda
Created November 8, 2019 09:54
Show Gist options
  • Save keisukefukuda/4f3103025363a33037873ee5d5db980a to your computer and use it in GitHub Desktop.
Save keisukefukuda/4f3103025363a33037873ee5d5db980a to your computer and use it in GitHub Desktop.
import chainermn
from chainer.datasets import get_cifar10
import pickle
comm = chainermn.create_communicator('pure_nccl')
if comm.rank == 0:
train, test = get_cifar10()
else:
train, test = None, None
try:
train = chainermn.scatter_dataset(train, comm, shuffle=True)
test = chainermn.scatter_dataset(test, comm, shuffle=True)
except (pickle.UnpicklingError, EOFError) as e:
print("**************************", flush=True)
print("Rank {} Error occured!!! ".format(comm.rank), flush=True)
print("{}: {}".format(type(e), str(e)), flush=True)
print("**************************", flush=True)
comm.mpi_comm.Abort(1)
else:
if comm.rank == 0:
print("OK")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment