High-level RNN MXNet Example¶

import os
import sys
import numpy as np
import mxnet as mx
from common.params_lstm import *
from common.utils import *

def create_symbol(CUDNN=True):
    # https://mxnet.incubator.apache.org/api/python/rnn.html
    data = mx.symbol.Variable('data')
    embedded_step = mx.symbol.Embedding(data=data, input_dim=MAXFEATURES, output_dim=EMBEDSIZE)
    
    # Fusing RNN layers across time step into one kernel
    # Improves speed but is less flexible
    # Currently only supported if using cuDNN on GPU
    if not CUDNN:
        gru_cell = mx.rnn.GRUCell(num_hidden=NUMHIDDEN)
    else:
        gru_cell = mx.rnn.FusedRNNCell(num_hidden=NUMHIDDEN, num_layers=1, mode='gru')
    
    begin_state = gru_cell.begin_state()
    # Call the cell to get the output of one time step for a batch.
    # TODO: TNC layout (sequence length, batch size, and feature dimensions) is faster for RNN
    outputs, states = gru_cell.unroll(length=MAXLEN, inputs=embedded_step, merge_outputs=False)
    
    fc1 = mx.symbol.FullyConnected(data=outputs[-1], num_hidden=2) 
    input_y = mx.symbol.Variable('softmax_label')  
    m = mx.symbol.SoftmaxOutput(data=fc1, label=input_y, name="softmax")
    return m

def init_model(m):
    if GPU:
        ctx = [mx.gpu(0)]
    else:
        ctx = mx.cpu()
    mod = mx.mod.Module(context=ctx, symbol=m)
    mod.bind(data_shapes=[('data', (BATCHSIZE, MAXLEN))],
             label_shapes=[('softmax_label', (BATCHSIZE, ))])
    # Glorot-uniform initializer
    mod.init_params(initializer=mx.init.Xavier(rnd_type='uniform'))
    mod.init_optimizer(optimizer='Adam', 
                       optimizer_params=(('learning_rate', LR),
                                         ('beta1', BETA_1),
                                         ('beta2', BETA_2),
                                         ('epsilon', EPS)))
    return mod