[FFI] Part3: npx.pooling, npx.dropout, npx.one_hot, npx.rnn

[barry-jin]

Description

part3 of #20096
Benchmarks:
npx.pooling

setup = """
from mxnet import np, npx
npx.set_np()
data = np.arange(9600).reshape(1, 1, 24, 20 ,20)
"""
stmt = """
out = npx.pooling(data=data, kernel=(2, 2, 2), stride=(2, 2, 2), \
                  pool_type='max', pooling_convention='valid', \
                  global_pool=False, cudnn_off=True, p_value=None, \
                  count_include_pad=None, layout=None)
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00015054992099999964
New FFI
9.832800099999961e-05

npx.dropout

setup = """
from mxnet import np, npx
npx.set_np()
inp = np.ones((10,10))
"""
stmt = """
out = npx.dropout(inp, p=0.5, mode="training", axes=None, cudnn_off=True)
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00011959457300000009
New FFI
7.979375999999982e-05

npx.one_hot

setup = """
from mxnet import np, npx
npx.set_np()
inp = np.zeros((10))
"""
stmt = """
out = npx.one_hot(inp, depth=2, on_value=1.0, off_value=0.0, dtype="float32")
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00012587712799999994
New FFI
6.992295699999973e-05

npx.rnn

setup = """
from mxnet import np, npx
npx.set_np()
seq_length = 10
batch_size = 4
input_size = 4
state_size = 1
data = np.random.normal(0, 1, (seq_length, batch_size, input_size))
state = np.random.normal(0, 1, (1, batch_size, state_size))
p = np.random.normal(0, 1, ((input_size + state_size + 2)*state_size),)
"""
stmt = """
out = npx.rnn(data=data, parameters=p, mode='rnn_tanh', \
              state=state, state_size=state_size, num_layers=1, \
              bidirectional=False, state_outputs=False, p=0.0, \
              use_sequence_length=False, projection_size=None, \
              lstm_state_clip_min=None,lstm_state_clip_max=None, \
              lstm_state_clip_nan=False)
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00017465537600000003
New FFI
8.456256000000017e-05

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

Changes

• Feature1, tests, (and when applicable, API doc)
• Feature2, tests, (and when applicable, API doc)

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

[barry-jin]

@mxnet-bot run ci [windows-cpu, windows-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [windows-cpu, windows-gpu]