lstm_test.py

import cupy as xp
import cupy.prof

import chainer
from chainer.functions.activation import lstm

B = 1024
D = 2048
x_shape = (B, D * 4)
c_shape = (B, D)

x = chainer.Variable(xp.random.uniform(-1, 1, x_shape).astype('f'))
c_p = chainer.Variable(xp.random.uniform(-1, 1, c_shape).astype('f'))
c = chainer.Variable(xp.random.uniform(-1, 1, c_shape).astype('f'))
gc = chainer.Variable(xp.random.uniform(-1, 1, c_shape).astype('f'))
gh= chainer.Variable(xp.random.uniform(-1, 1, c_shape).astype('f'))


with cupy.prof.time_range('lstm forward', color_id=0):
    c, h = lstm.LSTMGrad()(c_p, x, c, gc, gh)
    cupy.cuda.Stream.null.synchronize()

l = chainer.functions.sum(c) + chainer.functions.sum(h)
print(l)

with cupy.prof.time_range('lstm backward', color_id=0):
    l.backward()
    cupy.cuda.Stream.null.synchronize()

with_fuse.png

without_fuse.png