def test_using_device_with_name(device_instance1, device_instance2):
if device_instance1 == device_instance2:
return
device1 = device_instance1
device2 = device_instance2
chainerx.set_default_device(device1)
with chainerx.using_device(device2.name) as scope:
assert chainerx.get_default_device() == device2
assert scope.device is device2
with chainerx.using_device(device2.backend.name, device2.index) as scope:
assert chainerx.get_default_device() == device2
assert scope.device is device2
def test_backward_default_device(self):
# Default device in backward should be determined by arrays,
# otherwise, creation routines in backward do not create new arrays
# on the proper device.
device = chainerx.get_device('cuda:0')
shape = (2, 3)
dtype = numpy.float32
x1 = chainerx.full(shape, 3, dtype, device=device)
x2 = chainerx.full(shape, 5, dtype, device=device).require_grad()
backward_call_new_array = []
def backward_call_callback(call_arg):
backward_call_new_array.append(chainerx.empty(shape, dtype))
with chainerx.using_device('native:0'):
# forward
func = self.SimpleFunctionNode(backward_call_callback)
y1, y2 = func.apply((x1, x2))
# backward
y2.backward()
assert backward_call_new_array[0].device is device
def test_check_device_fail(shape, device, compare_device_spec):
dtype = 'float32'
a = chainerx.empty(shape, dtype, device=device)
with chainerx.using_device('native:1'):
with pytest.raises(AssertionError):
array_utils.check_device(a, compare_device_spec)
def test_check_device_fail(shape, device, compare_device_spec):
dtype = 'float32'
a = chainerx.empty(shape, dtype, device=device)
with chainerx.using_device('native:1'):
with pytest.raises(AssertionError):
array_utils.check_device(a, compare_device_spec)
def test_collective_obj_chx_gpu(self):
self.setup()
test_function_list = [
self.communicator.gather_obj, self.communicator.bcast_obj,
self.communicator.allreduce_obj
]
with chainerx.using_device("cuda"):
for func in test_function_list:
chx_array = chainerx.array([0])
with pytest.raises(ValueError):
func(chx_array)
chx_array_list = [[0], chainerx.array([1])]
with pytest.raises(ValueError):
func(chx_array_list)
chx_array_tuple = (0, chainerx.array([2]))
with pytest.raises(ValueError):
func(chx_array_tuple)
chx_array_dict_value = {0: chainerx.array([2])}
with pytest.raises(ValueError):
func(chx_array_dict_value)
chx_array_dict_key = {chainerx.array([2]): 0}
with pytest.raises(ValueError):
func(chx_array_dict_key)
chx_array_dict_set = {chainerx.array([2]), 0}
with pytest.raises(ValueError):
func(chx_array_dict_set)
self.teardown()
def test_send_obj_chx_gpu(self):
self.setup()
rank_next = (self.communicator.rank + 1) % self.communicator.size
with chainerx.using_device("cuda"):
chx_array = chainerx.array([0])
with pytest.raises(ValueError):
self.communicator.send_obj(chx_array, dest=rank_next)
chx_array_list = [[0], chainerx.array([1])]
with pytest.raises(ValueError):
self.communicator.send_obj(chx_array_list, dest=rank_next)
chx_array_tuple = (0, chainerx.array([2]))
with pytest.raises(ValueError):
self.communicator.send_obj(chx_array_tuple, dest=rank_next)
chx_array_dict_value = {0: chainerx.array([2])}
with pytest.raises(ValueError):
self.communicator.send_obj(chx_array_dict_value,
dest=rank_next)
chx_array_dict_key = {chainerx.array([2]): 0}
with pytest.raises(ValueError):
self.communicator.send_obj(chx_array_dict_key, dest=rank_next)
chx_array_dict_set = {chainerx.array([2]), 0}
with pytest.raises(ValueError):
self.communicator.send_obj(chx_array_dict_set, dest=rank_next)
self.teardown()
def test_backward_default_device(self):
# Default device in backward should be determined by arrays,
# otherwise, creation routines in backward do not create new arrays
# on the proper device.
device = chainerx.get_device('cuda:0')
shape = (2, 3)
dtype = numpy.float32
x1 = chainerx.full(shape, 3, dtype, device=device)
x2 = chainerx.full(shape, 5, dtype, device=device).require_grad()
backward_call_new_array = []
def backward_call_callback(call_arg):
backward_call_new_array.append(chainerx.empty(shape, dtype))
with chainerx.using_device('native:0'):
# forward
func = self.SimpleFunctionNode(backward_call_callback)
y1, y2 = func.apply((x1, x2))
# backward
y2.backward()
assert backward_call_new_array[0].device is device
def test_ascontiguousarray_from_chainerx_array_device():
with chainerx.using_device(chainerx.get_device('native:0')):
dev = chainerx.get_device('native:1') # Non default one
assert chainerx.get_default_device() is not dev
a = chainerx.arange(10, device=dev)
b = chainerx.ascontiguousarray(a)
assert b.is_contiguous is True
assert b.device is dev
def test_using_device(device_instance1, device_instance2):
if device_instance1 == device_instance2:
return
device1 = device_instance1
device2 = device_instance2
chainerx.set_default_device(device1)
with chainerx.using_device(device2) as scope:
assert chainerx.get_default_device() is device2
assert scope.device is device2
scope = chainerx.using_device(device2)
assert chainerx.get_default_device() == device1
assert scope.device is device2
with scope:
assert chainerx.get_default_device() == device2
assert scope.device is device2
assert chainerx.get_default_device() == device1
assert scope.device is device2
def test_send_recv_obj_chx_cpu(self):
self.setup()
with chainerx.using_device("native"):
chx_array = chainerx.array([0])
self.check_send_recv_obj(chx_array)
chx_array = chainerx.array([1])
self.check_send_recv_obj(chx_array, tag=1)
chx_array = chainerx.array([2])
self.check_send_recv_obj(chx_array, tag=2, use_any_recv=False)
self.teardown()
def test_observation_aggregator_gpu_chainerx(use_chainer_variable,
communicate_interval):
xp = chainerx
communicator = chainermn.create_communicator('pure_nccl')
device = get_device(communicator.intra_rank, True)
with chainerx.using_device(device.device):
if use_chainer_variable:
run_test_observation_aggregator(communicator,
xp,
use_chainer_variable,
communicate_interval,
use_gpu=True)
else:
with pytest.raises(ValueError):
run_test_observation_aggregator(communicator,
xp,
use_chainer_variable,
communicate_interval,
use_gpu=True)
def test_init_chainerx_with_default_device(self):
device = chainerx.get_device('native:1')
with chainerx.using_device(device):
array = self._generate_array(chainerx, 'float64')
assert array.device is device
def test_init_chainerx_with_default_device(self):
device = chainerx.get_device('native:1')
with chainerx.using_device(device):
array = self._generate_array(chainerx, 'float64')
assert array.device is device
def main():
parser = argparse.ArgumentParser(description='Compare chainer vs chainerx')
parser.add_argument('--batchsize', '-b', type=int, default=100)
parser.add_argument('--epoch', '-e', type=int, default=10)
parser.add_argument('--gpu', '-g', type=int, default=0, choices=[-1, 0, 1, 2, 3])
parser.add_argument('--chxon', '-c', type=int, default=1)
args = parser.parse_args()
# setup
start = time.time()
chx.available = True if args.chxon == 1 else False
batch_size = args.batchsize
# get MNIST
train, test = chainer.datasets.get_mnist()
if chx_available == True:
device_name = 'cuda:{}'.format(args.gpu)
# data
with chx.using_device(device_name):
train_images, train_labels = map(lamda d:chx.asarray(d), train._datasets)
test_images, test_labels = map(lamda d:chx.asarray(d), test._datasets)
# model
chx.set_default_device(device_name)
model = MLP(n_units=1000, n_out=10)
optimizer = SGD(lr=0.01)
else:
device_name = args.gpu
# data
train_iter = chainer.iterators.SerialIterator(train, batch_size)
test_iter = chainer.iterators.SerialIterator(train, batch_size, repeat=False, shuffle=False)
# model
model = MLP_chain(n_units=1000, n_out=10)
model.to_gpu()
chainer.cuda.get_device_from_id(device_name).use()
optimizer = chainer.optimizers.SGD(lr=0.01)
optimizer.setup(model)
N_train, N_test = len(train), len(test)
all_indices_np = np.arange(N_train, dtype=np.int64) # for chainerx
epoch = 0
while epoch <= args.epoch:
epoch += 1
if chx_available == True:
np.random.shuffle(all_indices_np)
all_indices = chx.array(all_indices_np)
for i in range(0, N_train, batch_size):
# time 1
if chx_available == True:
indices = all_indices[i: i + batch_size]
x = train_images.take(indices, axis=0)
t = train_labels.take(indices, axis=0)
else:
batch = train_iter.next()
x, t = convert.concat_examples(batch, device=device_name)
y = model.forward(x) # time 2
# time 3
if chx_available == True:
loss = compute_loss(y, t)
else:
loss = F.softmax_cross_entropy(y, t)
model.cleargrads()
loss.backward() # time 4
optimizer.update() # time 5
if chx_available == True:
with chx.no_backprop_mode():
total_loss = chx.array(0, dtype=chx.float32)
num_correct = chx.array(0, dtype=chx.int64)
for i in range(0, N_test, batch_size):
x = test_images[i:min(i + batch_size, N_test)]
x = test_labels[i:min(i + batch_size, N_test)]
y = model.forward(x)
total_loss += compute_loss(y, t) * len(t)
num_correct += (y.argmax(axis=1).astype(t.dtype) == t).astype(chx.int32).sum()
else:
test_iter.reset()
with chainer.using_config('enable_backprop', False):
total_loss = 0
num_correct = 0
for batch in test_iter:
x, t = convert.concat_examples(batch, device=device_name)
y = model.forward(x)
total_loss += float(F.softmax_cross_entropy(y, t).array) * len(t)
num_correct += float(F.accuracy(y, t).array) * len(t)
mean_loss = float(total_loss) / N_test
accuracy = int(num_correct) / N_test
elapsed_time = time.time() - start
print('epoch {} ... loss={}, accuracy, elapsed_time={}'.format(
epoch, mean_loss, accuracy, elapsed_time))
def test_use(self, backend_config):
device = chainer.get_device(backend_config.chainerx_device)
with chainerx.using_device('native:1'):
device.use()
assert device.device is chainerx.get_default_device()
def create_context(self):
# Returns a context that sets the default device.
return chainerx.using_device(self.device)
def test_use(self, backend_config):
device = chainer.get_device(backend_config.chainerx_device)
with chainerx.using_device('native:1'):
device.use()
assert device.device is chainerx.get_default_device()
def create_context(self):
# Returns a context that sets the default device.
return chainerx.using_device(self.device)
def test_early_stopping_trigger_with_accuracy_gpu_chx(self):
self.communicator = chainermn.create_communicator('pure_nccl')
self.xp = chainerx
chainermn.testing.get_device(self.communicator.intra_rank, True).use()
with chainerx.using_device("cuda", self.communicator.intra_rank):
self.run_test_early_stopping_trigger_with_accuracy()
