Improve backward performance
@@ -351,6 +351,7 @@ Actual: {0}'''.format(type(data))
"""
if self.creator is None:
return
+ initial_device = cuda.Device()
is_debug = chainer.is_debug()
@@ -385,18 +386,21 @@ Actual: {0}'''.format(type(data))
if func._n_local_function_hooks != 0:
hooks = collections.OrderedDict(hooks)
hooks.update(func.local_function_hooks)
+
+ cuda.get_device(*(in_data + out_grad)).use()
for hook in six.itervalues(hooks):
hook.backward_preprocess(func, in_data, out_grad)
- with cuda.get_device(*(in_data + out_grad)):
gxs = func.backward(in_data, out_grad)
assert len(gxs) == len(in_data)
for hook in six.itervalues(hooks):
hook.backward_postprocess(func, in_data, out_grad)
if is_debug:
- if any(gx is not None and
- cuda.get_array_module(gx).isnan(gx).any()
- for gx in gxs):
+ for gx in gxs:
+ if gx is None:
+ continue
+ cuda.get_device(gx).use()
+ if cuda.get_array_module(gx).isnan(gx).any():
msg = 'NaN is detected on backward computation'
raise RuntimeError(msg)
@@ -412,13 +416,14 @@ Actual: {0}'''.format(type(data))
# Accumulate the gradient to x. It is a bit tricky to handle
# branches and parameter gradient accumulation correctly.
- with cuda.get_device(gx):
id_x = id(x)
if x.creator is None: # leaf
if x._grad is None:
x.grad = gx
need_copy.add(id_x)
- elif id_x in need_copy:
+ else:
+ cuda.get_device(gx).use()
+ if id_x in need_copy:
x.grad = utils.force_array(x.grad + gx) # copy
need_copy.remove(id_x)
else:
@@ -429,12 +434,15 @@ Actual: {0}'''.format(type(data))
x.grad = gx
seen_vars.add(id_x)
need_copy.add(id_x)
- elif id_x in need_copy: # 2nd visit
+ else:
+ cuda.get_device(gx).use()
+ if id_x in need_copy: # 2nd visit
x._grad = utils.force_array(gx + x._grad) # copied
need_copy.remove(id_x)
else: # 3rd or later visit
x._grad += gx
del gxs # to reduce memory usage
+ initial_device.use()
def unchain_backward(self):
"""Deletes references between variables and functions backward.
