Forward-only rnns

Aapo Kyrola · Yangqing Jia · commit 9cb901caf0d1 · 2017-04-24T15:52:27.000-07:00
Summary:
Added option to recurrent_net and RNNCell's for forward_only. If this is set, the backward_step_net is not passed to the operator.
When backward_step_net is not available, operator knows it is in forward_only mode and does not create workspaces for each step but cycles
through only one private workspace.

Note: we could avoid doing a lot of work in recurrent.py:recurrent_network call when backward step is not needed, but doing that nicely requires
more refactoring that I did not want to do now. Thus, we create the backward step nets etc, but just don't pass it to the op.

This can be used to create more efficient inference models. You can also sanitize existing inference nets and remove the backward_step_net argument to
get the benefits.

Reviewed By: salexspb

Differential Revision: D4916482

fbshipit-source-id: c99b93c9cb897c32b0f449253f7f6d6a942618ad
diff --git a/caffe2/operators/recurrent_network_op.h b/caffe2/operators/recurrent_network_op.h
@@ -257,6 +257,13 @@ class RecurrentNetworkOp final : public Operator<Context> {
           ri, seqLen, batchSize, sharedWs_, &context_);
     }
 
+    // If we don't have a backward step net, this operator is forward_only
+    // and we can avoid creating multiple workspaces.
+
+    bool has_backward_pass =
+        OperatorBase::GetSingleArgument<string>("backward_step_net", "") != "";
+
+    // With backward pass: we need to create workspace for each timestep
     detail::ScratchWorkspaces* scratch =
         OperatorBase::Output<detail::ScratchWorkspaces>(OutputSize() - 1);
     std::vector<std::shared_ptr<Workspace>>& stepWorkspaces =
@@ -271,13 +278,15 @@ class RecurrentNetworkOp final : public Operator<Context> {
     // have to be stored in step workspaces but can be shared.
     initializeBlobsToRecomputeOnBackward(forwardSharedWs.get());
 
-    if (seqLen > stepWorkspaces.size()) {
+    if (has_backward_pass && seqLen > stepWorkspaces.size()) {
       stepWorkspaces.resize(seqLen);
     }
 
     for (auto t = 0; t < seqLen; ++t) {
-      auto& currentStepWorkspace = stepWorkspaces[t];
+      auto& currentStepWorkspace =
+          (has_backward_pass ? stepWorkspaces[t] : forwardSharedWs);
       if (!currentStepWorkspace) {
+        CHECK(has_backward_pass);
         currentStepWorkspace =
             std::make_shared<Workspace>(forwardSharedWs.get());
       }
diff --git a/caffe2/python/lstm_benchmark.py b/caffe2/python/lstm_benchmark.py
@@ -80,6 +80,7 @@ def create_model(args, queue, label_queue, input_shape):
             dim_out=args.hidden_dim,
             scope="lstm1",
             memory_optimization=args.memory_optimization,
+            forward_only=args.forward_only,
         )
     elif args.implementation == "cudnn":
         # We need to feed a placeholder input so that RecurrentInitOp
@@ -104,7 +105,8 @@ def create_model(args, queue, label_queue, input_shape):
         ['softmax', 'loss'],
     )
 
-    model.AddGradientOperators([loss])
+    if not args.forward_only:
+        model.AddGradientOperators([loss])
 
     # carry states over
     model.net.Copy(last_hidden, hidden_init)
@@ -232,6 +234,11 @@ def GetArgumentParser():
         action="store_true",
         help="Whether to use memory optimized LSTM or not",
     )
+    parser.add_argument(
+        "--forward_only",
+        action="store_true",
+        help="Whether to run only forward pass"
+    )
 
     return parser
 
diff --git a/caffe2/python/operator_test/rnn_cell_test.py b/caffe2/python/operator_test/rnn_cell_test.py
@@ -378,10 +378,10 @@ def test_lstm_unit_recurrent_network(self, n, d, t, dc, gc):
                 gc, op, inputs, i, [0, 1],
                 input_device_options=input_device_options)
 
-
     @given(
         input_tensor=lstm_input(),
         forget_bias=st.floats(-10.0, 10.0),
+        fwd_only=st.booleans(),
     )
     @ht_settings(max_examples=25)
     def test_lstm_main(self, **kwargs):
@@ -393,7 +393,7 @@ def test_lstm_main(self, **kwargs):
                                    **kwargs)
 
     def lstm_base(self, lstm_type, outputs_with_grads, memory_optim,
-                  input_tensor, forget_bias):
+                  input_tensor, forget_bias, fwd_only):
         print("LSTM test parameters: ", locals())
         create_lstm, ref = lstm_type
         t, n, d = input_tensor.shape
@@ -412,7 +412,8 @@ def lstm_base(self, lstm_type, outputs_with_grads, memory_optim,
             d, d, scope="external/recurrent",
             outputs_with_grads=outputs_with_grads,
             memory_optimization=memory_optim,
-            forget_bias=forget_bias)
+            forget_bias=forget_bias,
+            forward_only=fwd_only)
 
         op = model.net._net.op[-1]
 
@@ -447,16 +448,17 @@ def generate_random_state(n, d):
         )
 
         # Checking for input, gates_t_w and gates_t_b gradients
-        for param in range(5):
-            self.assertGradientChecks(
-                device_option=hu.cpu_do,
-                op=op,
-                inputs=inputs,
-                outputs_to_check=param,
-                outputs_with_grads=outputs_with_grads,
-                threshold=0.01,
-                stepsize=0.005,
-            )
+        if not fwd_only:
+            for param in range(5):
+                self.assertGradientChecks(
+                    device_option=hu.cpu_do,
+                    op=op,
+                    inputs=inputs,
+                    outputs_to_check=param,
+                    outputs_with_grads=outputs_with_grads,
+                    threshold=0.01,
+                    stepsize=0.005,
+                )
 
     @given(encoder_output_length=st.integers(1, 3),
            encoder_output_dim=st.integers(1, 3),
diff --git a/caffe2/python/recurrent.py b/caffe2/python/recurrent.py
@@ -8,10 +8,11 @@
 from caffe2.python import core
 from caffe2.python.scope import CurrentNameScope
 
+
 def recurrent_net(
         net, cell_net, inputs, initial_cell_inputs,
         links, timestep=None, scope=None, outputs_with_grads=(0,),
-        recompute_blobs_on_backward=None,
+        recompute_blobs_on_backward=None, forward_only=False,
 ):
     '''
     net: the main net operator should be added to
@@ -43,6 +44,8 @@ def recurrent_net(
     recompute_blobs_on_backward: specify a list of blobs that will be
                  recomputed for backward pass, and thus need not to be
                  stored for each forward timestep.
+
+    forward_only: if True, only forward steps are executed
     '''
     assert len(inputs) == 1, "Only one input blob is supported so far"
 
@@ -77,54 +80,55 @@ def s(name):
     inner_outputs = list(cell_net.Proto().external_output)
     # These gradients are expected to be available during the backward pass
     inner_outputs_map = {o: o + '_grad' for o in inner_outputs}
+    recompute_blobs_on_backward = set()
 
     # compute the backward pass of the cell net
-    backward_ops, backward_mapping = core.GradientRegistry.GetBackwardPass(
-        cell_net.Proto().op, inner_outputs_map)
-    backward_mapping = {str(k): v for k, v in backward_mapping.items()}
-    backward_cell_net = core.Net("RecurrentBackwardStep")
-    del backward_cell_net.Proto().op[:]
-
-    if recompute_blobs_on_backward is not None:
-        # Insert operators to re-compute the specified blobs.
-        # They are added in the same order as for the forward pass, thus
-        # the order is correct.
-        recompute_blobs_on_backward = set(
-            [str(b) for b in recompute_blobs_on_backward]
-        )
-        for op in cell_net.Proto().op:
-            if not recompute_blobs_on_backward.isdisjoint(set(op.output)):
-                backward_cell_net.Proto().op.extend([op])
-                assert set(op.output).issubset(recompute_blobs_on_backward), \
-                       'Outputs {} are output by op but not recomputed: {}'.format(
-                            set(op.output) - recompute_blobs_on_backward,
-                            op
-                       )
+    if not forward_only:
+        backward_ops, backward_mapping = core.GradientRegistry.GetBackwardPass(
+            cell_net.Proto().op, inner_outputs_map)
+        backward_mapping = {str(k): v for k, v in backward_mapping.items()}
+
+        backward_cell_net = core.Net("RecurrentBackwardStep")
+        del backward_cell_net.Proto().op[:]
+
+        if recompute_blobs_on_backward is not None:
+            # Insert operators to re-compute the specified blobs.
+            # They are added in the same order as for the forward pass, thus
+            # the order is correct.
+            recompute_blobs_on_backward = {str(b) for b in
+                                           recompute_blobs_on_backward}
+
+            for op in cell_net.Proto().op:
+                if not recompute_blobs_on_backward.isdisjoint(set(op.output)):
+                    backward_cell_net.Proto().op.extend([op])
+                    # This fires if other outputs than the declared
+                    # are computed by the ops that are recomputed
+                    assert set(op.output).issubset(recompute_blobs_on_backward)
+
+        backward_cell_net.Proto().op.extend(backward_ops)
+        # compute blobs used but not defined in the backward pass
+        backward_ssa, backward_blob_versions = core.get_ssa(
+            backward_cell_net.Proto())
+        undefined = core.get_undefined_blobs(backward_ssa)
+
+        # also add to the output list the intermediate outputs of fwd_step that
+        # are used by backward.
+        ssa, blob_versions = core.get_ssa(cell_net.Proto())
+        scratches = [
+            blob for (blob, ver) in blob_versions.items()
+            if ver > 0 and
+            blob in undefined and
+            blob not in cell_net.Proto().external_output]
+        backward_cell_net.Proto().external_input.extend(scratches)
+        backward_cell_net.Proto().type = 'simple'
     else:
-        recompute_blobs_on_backward = set()
-
-    backward_cell_net.Proto().op.extend(backward_ops)
-    # compute blobs used but not defined in the backward pass
-    backward_ssa, backward_blob_versions = core.get_ssa(
-        backward_cell_net.Proto())
-    undefined = core.get_undefined_blobs(backward_ssa)
-
-    # also add to the output list the intermediate outputs of fwd_step that
-    # are used by backward.
-    ssa, blob_versions = core.get_ssa(cell_net.Proto())
-    scratches = [
-        blob for (blob, ver) in blob_versions.items()
-        if ver > 0 and
-        blob in undefined and
-        blob not in cell_net.Proto().external_output]
-    backward_cell_net.Proto().external_input.extend(scratches)
+        backward_cell_net = None
 
     all_inputs = [i[1] for i in inputs] + [
         x[1] for x in initial_cell_inputs] + references
     all_outputs = []
 
     cell_net.Proto().type = 'simple'
-    backward_cell_net.Proto().type = 'simple'
 
     # Internal arguments used by RecurrentNetwork operator
 
@@ -153,53 +157,58 @@ def s(name):
         cell_output = links[str(cell_input)]
         forward_links.append((cell_input, state, 0))
         forward_links.append((cell_output, state, 1))
-        backward_links.append((cell_output + "_grad", states_grad, 1))
 
-        backward_cell_net.Proto().external_input.append(
-            str(cell_output) + "_grad")
         aliases.append((state, cell_output + "_all", 1))
         aliases.append((state, cell_output + "_last", -1))
         all_outputs.extend([cell_output + "_all", cell_output + "_last"])
 
         recurrent_states.append(state)
 
-        recurrent_input_grad = cell_input + "_grad"
-        if not backward_blob_versions.get(recurrent_input_grad, 0):
-            # If nobody writes to this recurrent input gradient, we need
-            # to make sure it gets to the states grad blob after all.
-            # We do this by using backward_links which triggers an alias
-            # This logic is being used for example in a SumOp case
-            backward_links.append(
-                (backward_mapping[cell_input], states_grad, 0))
-        else:
-            backward_links.append((cell_input + "_grad", states_grad, 0))
-
-    for reference in references:
-        # Similar to above, in a case of a SumOp we need to write our parameter
-        # gradient to an external blob. In this case we can be sure that
-        # reference + "_grad" is a correct parameter name as we know how
-        # RecurrentNetworkOp gradient schema looks like.
-        reference_grad = reference + "_grad"
-        if (reference in backward_mapping and
-                reference_grad != str(backward_mapping[reference])):
-            # We can use an Alias because after each timestep
-            # RNN op adds value from reference_grad into and _acc blob
-            # which accumulates gradients for corresponding parameter accross
-            # timesteps. Then in the end of RNN op these two are being
-            # swaped and reference_grad blob becomes a real blob instead of
-            # being an alias
-            backward_cell_net.Alias(
-                backward_mapping[reference], reference_grad)
+        if backward_cell_net is not None:
+            backward_links.append((cell_output + "_grad", states_grad, 1))
+            backward_cell_net.Proto().external_input.append(
+                str(cell_output) + "_grad")
+
+            recurrent_input_grad = cell_input + "_grad"
+            if not backward_blob_versions.get(recurrent_input_grad, 0):
+                # If nobody writes to this recurrent input gradient, we need
+                # to make sure it gets to the states grad blob after all.
+                # We do this by using backward_links which triggers an alias
+                # This logic is being used for example in a SumOp case
+                backward_links.append(
+                    (backward_mapping[cell_input], states_grad, 0))
+            else:
+                backward_links.append((cell_input + "_grad", states_grad, 0))
 
     for input_t, input_blob in inputs:
         forward_links.append((str(input_t), str(input_blob), 0))
-        backward_links.append((
-            backward_mapping[str(input_t)], str(input_blob) + "_grad", 0
-        ))
-    backward_cell_net.Proto().external_input.extend(
-        cell_net.Proto().external_input)
-    backward_cell_net.Proto().external_input.extend(
-        cell_net.Proto().external_output)
+
+    if backward_cell_net is not None:
+        for reference in references:
+            # Similar to above, in a case of a SumOp we need to write our parameter
+            # gradient to an external blob. In this case we can be sure that
+            # reference + "_grad" is a correct parameter name as we know how
+            # RecurrentNetworkOp gradient schema looks like.
+            reference_grad = reference + "_grad"
+            if (reference in backward_mapping and
+                    reference_grad != str(backward_mapping[reference])):
+                # We can use an Alias because after each timestep
+                # RNN op adds value from reference_grad into and _acc blob
+                # which accumulates gradients for corresponding parameter accross
+                # timesteps. Then in the end of RNN op these two are being
+                # swaped and reference_grad blob becomes a real blob instead of
+                # being an alias
+                backward_cell_net.Alias(
+                    backward_mapping[reference], reference_grad)
+
+        for input_t, input_blob in inputs:
+            backward_links.append((
+                backward_mapping[str(input_t)], str(input_blob) + "_grad", 0
+            ))
+        backward_cell_net.Proto().external_input.extend(
+            cell_net.Proto().external_input)
+        backward_cell_net.Proto().external_input.extend(
+            cell_net.Proto().external_output)
 
     def unpack_triple(x):
         if x:
@@ -210,18 +219,28 @@ def unpack_triple(x):
     # Splitting to separate lists so we can pass them to c++
     # where we ensemle them back
     link_internal, link_external, link_offset = unpack_triple(forward_links)
-    backward_link_internal, backward_link_external, backward_link_offset = \
-        unpack_triple(backward_links)
     alias_src, alias_dst, alias_offset = unpack_triple(aliases)
 
-    params = [x for x in references if x in backward_mapping.keys()]
     recurrent_inputs = [str(x[1]) for x in initial_cell_inputs]
 
-    global _workspace_seq
+    backward_args = {}
+    if backward_cell_net is not None:
+        backward_link_internal, backward_link_external, backward_link_offset = \
+            unpack_triple(backward_links)
+        params = [x for x in references if x in backward_mapping.keys()]
+        backward_args = {
+            'param': map(all_inputs.index, params),
+            'backward_link_internal': map(str, backward_link_internal),
+            'backward_link_external': map(str, backward_link_external),
+            'backward_link_offset': backward_link_offset,
+            'backward_step_net': str(backward_cell_net.Proto()),
+            'outputs_with_grads': outputs_with_grads,
+            'recompute_blobs_on_backward': map(str, recompute_blobs_on_backward)
+        }
+
     results = net.RecurrentNetwork(
         all_inputs,
         all_outputs + [s("step_workspaces")],
-        param=map(all_inputs.index, params),
         alias_src=alias_src,
         alias_dst=map(str, alias_dst),
         alias_offset=alias_offset,
@@ -230,14 +249,9 @@ def unpack_triple(x):
         link_internal=map(str, link_internal),
         link_external=map(str, link_external),
         link_offset=link_offset,
-        backward_link_internal=map(str, backward_link_internal),
-        backward_link_external=map(str, backward_link_external),
-        backward_link_offset=backward_link_offset,
         step_net=str(cell_net.Proto()),
-        backward_step_net=str(backward_cell_net.Proto()),
         timestep="timestep" if timestep is None else str(timestep),
-        outputs_with_grads=outputs_with_grads,
-        recompute_blobs_on_backward=map(str, recompute_blobs_on_backward)
+        **backward_args
     )
     # The last output is a list of step workspaces,
     # which is only needed internally for gradient propogation
diff --git a/caffe2/python/rnn_cell.py b/caffe2/python/rnn_cell.py
