rnn training example?

[mingfeima]

is there going to be a rnn training example? say how to allocate the workspace shared between rnn forward and rnn backward?

[emfomenk]

Hi @mingfeima,

Workspace allocation and sharing between forward and backward passes are common for the entire MKL-DNN API. The typical steps are:

1. create forward primitive descriptor (note that prop_kind should be equal to forward_training)
2. query this primitive descriptor for workspace memory primitive descriptor
3. if that workspace_pd is not NULL, query its size via mkldnn_memory_primitive_desc_get_size()
4. allocate enough memory for the workspace, create workspace memory and set handle to allocated memory
5. pass workspace as one of the outputs for forward primitive
6. on backward pass put the same workspace as one of the inputs

The corresponding places can be found:

• C example: lines 436 .. 450
• C++ example: lines 215 .. 220

[mingfeima]

@emfomenk thanks for the detailed reply!

Is mkldnn rnn going to support dropout, like cuDNN?
And does mkldnn has GRU support at the moment?

[emfomenk]

HI @mingfeima,

GRU is something we are currently working on. Do you have any comments in what particular implementation you are interested in? There are a bunch of different incompatible implementations (TF, MxNet, PaddlePaddle, cuDNN, ...).

Alas, dropout is not supported. Could you please elaborate on where is it used? Or the ask is more about compatibility?

[mingfeima]

@emfomenk i am working on PyTorch in which GRU is implemented in the following way, should be identical to cuDNN version. I wasn't aware of other implementations so i can't make compare.

    gi = F.linear(input, w_ih, b_ih)
    gh = F.linear(hidden, w_hh, b_hh)
    i_r, i_i, i_n = gi.chunk(3, 1)
    h_r, h_i, h_n = gh.chunk(3, 1)

    resetgate = F.sigmoid(i_r + h_r)
    inputgate = F.sigmoid(i_i + h_i)
    newgate = F.tanh(i_n + resetgate * h_n)
    hy = newgate + inputgate * (hidden - newgate)

And yes, it's more about compatibility at the moment from my side.

[mratsim]

I've just started implementing GRU in my own deep learning framework and did a comprehensive overview of all the implementations in the wild here, including a reference to the heavy optimizations proposed by Baidu Research.

Here are the equations I've found:

• ◦ is element-wise multiplication
• * is tensor multiplication
  Tensorflow

  r = sigmoid(W_{ir} * x + b_{ir} + W_{hr} * h + b_{hr}) 
  z = sigmoid(W_{iz} * x + b_{iz} + W_{hz} * h + b_{hz})
  n = tanh(W_{in} * x + b_{in} +  W_{hn} * (r ◦ h) + b_{hn}))
  h' = (1 - z) ◦ n + z ◦ h
  

  PyTorch equations

  r = sigmoid(W_{ir} * x + b_{ir} + W_{hr} * h + b_{hr}) 
  z = sigmoid(W_{iz} * x + b_{iz} + W_{hz} * h + b_{hz})
  n = tanh(W_{in} * x + b_{in} + r ◦ (W_{hn} * h + b_{hn}))
  h' = (1 - z) ◦ n + z ◦ h
  

  Note that in the paper equations are:

  r = sigmoid(W_{ir} * x + b_{ir} + W_{hr} * h + b_{hr}) 
  z = sigmoid(W_{iz} * x + b_{iz} + W_{hz} * h + b_{hz})
  n = tanh(W_{in} * x + b_{in} +  W_{hn} * (r ◦ h) + b_{hn}))
  h' = (1 - z) ◦ h + z ◦ n
  

  And CuDNN

  it = σ(Wi * xt + Ri * ht-1 + bWi + bRu)
  rt = σ(Wr * xt + Rr * ht-1 + bWr + bRr)
  h't = tanh(Wh * xt + rt ◦ (Rh * ht-1 + bRh) + bWh)
  ht = (1 - it)◦h't + it◦ht-1
  

Note that PyTorch/CuDNN have the same equations which allows reusing weights. Reusing weights between CPU and GPU is a pain point in Keras at the moment:

• Keras: weights on GPU cannot be reused on CPU and solutions (i.e. redoing a CPU layer):
  • LSTM
  • GRU

My own implementation in Nim outputs the same result as PyTorch but limits the number of temporary allocations by reusing buffers and leveraging loopfusion (sigmoid(a+b) or tanh(x + y ◦ z) in single pass.

proc gru_cell_forward*[T: SomeReal](input, hidden,
                                    w_input, w_recur,
                                    b_input, b_recur: Tensor[T],
                                    result: var Tensor[T]) =
  let
    H = hidden.shape[1]
    # Slices
    sr = (0 ..< H)|1
    sz = (H ..< 2*H)|1
    srz = (0 ..< 2*H)|1
    sn = (2*H ..< 3*H)|1

  var Wx, Rh: Tensor[T] # TODO, pass those as parameter to allow buffer reuse
  # Resulting shape [batch_size, 3*H]
  linear(input, w_input, b_input, Wx)
  linear(hidden, w_recur, b_recur, Rh)

  # To reduce allocations, we compute reset gate r
  # and update gate z in the previous buffers
  # We keep them concatenated to improve throughput
  var reset_update = Wx[_, srz] # shape [batch_size, 2*H]
  apply2_inline(reset_update, Rh[_, srz]):
    sigmoid(x + y)

  # We also reuse the previous buffer for the candidate hidden state n
  # shape [H, batch_size]
  var n = Wx[_, sn] # shape [batch_size, H]
  apply3_inline(n, reset_update[_, sr], Rh[_, sn]):
    tanh(x + y * z)

  # Compute the next hidden state
  result = map3_inline(Wx[_, sz], n, hidden):
    (1 - x) * y + x * z

[vpirogov]

@mratsim, thanks for sharing your findings and insights on the weight reuse problem. @igorsafo, @irinasok is this consistent with the picture we have on GRU flavors?

[irinasok]

Hi @mingfeima, @mratsim!

Thanks for detailed overview! We are looking into all of the implementations. Right now we are preparing original version of GRU that corresponds to Tensorflow implementation:

r = sigmoid(W_{ir} * x + b_{ir} + W_{hr} * h + b_{hr}) 
z = sigmoid(W_{iz} * x + b_{iz} + W_{hz} * h + b_{hz})
n = tanh(W_{in} * x + b_{in} +  W_{hn} * (r ◦ h) + b_{hn}))
h' = (1 - z) ◦ n + z ◦ h

But we also have plans to implement CuDNN variant and leave a switch for user to choose between those two.

[vpirogov]

Initial GRU support is added in 0282e6b. Let us know what you think.

[mingfeima]

@vpirogov as long as cuDNN flavor of GRU is provided, i am fine with current interface.
@emfomenk got a few questions when integrating MKLDNN RNN into pytorch.

1. how MKLDNN differs bias_ih and bias_hh? and also grad_bias_ih and grad_bias_hh? In the simple_rnn example, there is only one bias described by mkldnn::memory::format::ldgo
2. which is the memory::format for hx, cx, hy, cy,the following is pytorch shape
   hx - hidden state input, shape (num_layers * num_directions, batch, hidden_size)
   cx - cell state input, shape (num_layers * num_directions, batch, hidden_size)
   hy - hidden state output, shape (num_layers * num_directions, batch, hidden_size)
   cy - cell state output, shape (num_layers * num_directions, batch, hidden_size)
   is it mkldnn_ldsnc? if so, num_states refers to hx and cx for LSTM, hx for GRU?
   then i need to cancat hx and cx into a 5D tensor along num_states?
3. Do you implement layer fusion? basically i need to know if there is performance penalty in case i break a multi-layer RNN primitive into separate layer primitives.
4. Do you have packed_sequence input at the moment? Right now i am unpacking the packed input myself and shipping input to MKLDNN in TNC format.