Skip to content

Instantly share code, notes, and snippets.

@mitmul
Last active May 7, 2018 06:38
Show Gist options
  • Save mitmul/571a155d651df3b0dbe6ec129c846cb1 to your computer and use it in GitHub Desktop.
Save mitmul/571a155d651df3b0dbe6ec129c846cb1 to your computer and use it in GitHub Desktop.
docker build -t mitmul/repro:cupy-issue-1222 .
nvidia-docker run --rm \
-v $PWD:/root \
-ti mitmul/repro:cupy-issue-1222 \
python3 train.py
FROM nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04
RUN apt-get update && apt-get install -y \
python3 \
python3-dev \
python3-dbg \
python3-pip \
python3-wheel \
git \
wget \
curl \
vim
RUN pip3 install \
cython \
chainer==4.0.0 \
cupy-cuda90==4.0.0
WORKDIR /root
import chainer
import chainer.functions as F
import chainer.links as L
from chainer import training
class MLP(chainer.Chain):
def __init__(self, n_units, n_out):
super(MLP, self).__init__()
with self.init_scope():
# the size of the inputs to each layer will be inferred
self.l1 = L.Linear(None, n_units) # n_in -> n_units
self.l2 = L.Linear(None, n_units) # n_units -> n_units
self.l3 = L.Linear(None, n_out) # n_units -> n_out
def __call__(self, x):
h1 = F.relu(self.l1(x))
h2 = F.relu(self.l2(h1))
return self.l3(h2)
def train():
train, test = chainer.datasets.get_mnist()
batch_size = 64
learning_rate = 0.05
model = L.Classifier(MLP(1000, 10))
optimizer = chainer.optimizers.MomentumSGD(learning_rate)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(5e-4))
# Set up a trainer
num_gpus = 2
devices = range(num_gpus)
# this is just to force the error
chainer.cuda.get_device_from_id(0).use()
train_iters = [chainer.iterators.MultiprocessIterator(i, batch_size, n_processes=num_gpus) \
for i in chainer.datasets.split_dataset_n_random(train, len(devices))]
updater = training.updaters.MultiprocessParallelUpdater(train_iters, optimizer, devices=range(num_gpus))
updater.setup_workers()
if __name__=="__main__":
train()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment