#########################################################
#
# DO NOT EDIT THIS FILE. IT IS GENERATED AUTOMATICALLY. #
# PLEASE LOOK INTO THE README FOR MORE INFORMATION.     #
#
#########################################################


# coding: utf-8

# # Control Ops Tutorial
# 
# In this tutorial we show how to use control flow operators in Caffe2 and give some details on their underlying implementations.

# ### Conditional Execution

# Let's start with conditional operator. We will demostrate how to use it in two Caffe2 APIs used for building nets: NetBuilder and Brew.

# In the first example, we first define several blobs and then use 'If' operator to set value of one of them conditionally depending on values of other blobs.

# In[1]:


from caffe2.python import workspace
from caffe2.python.core import Plan, to_execution_step, Net
from caffe2.python.net_builder import ops, NetBuilder


# In[2]:


with NetBuilder() as nb:
    ops.Const(0.0, blob_out="zero")
    ops.Const(1.0, blob_out="one")
    ops.Const(0.5, blob_out="x")
    ops.Const(0.0, blob_out="y")
    with ops.IfNet(ops.GT(["x", "zero"])):
        ops.Copy("one", "y")
    with ops.Else():
        ops.Copy("zero", "y")


# Note the usage of NetBuilder's ops.IfNet and ops.Else calls: ops.IfNet accepts a blob reference or blob name as an input, it expects an input blob to have a scalar value convertible to bool, also note that optional ops.Else is at the same level as ops.IfNet and immediately follows corresponding ops.IfNet. Let's execute resulting net (execution step) and check values of blobs.

# In[3]:


plan = Plan('if_net_test')
plan.AddStep(to_execution_step(nb))
ws = workspace.C.Workspace()
ws.run(plan)
print('x = ', ws.blobs["x"].fetch())
print('y = ', ws.blobs["y"].fetch())


# Before going further, it's important to understand the semantics of execution blocks ('then' and 'else' branches in the example above), i.e. handling of reads and writes into global (defined outside of the block) and local (defined inside the block) blobs.

# NetBuilder's uses the following set of rules:
#  - In NetBuilder's syntax, blob's declaration and definition occur at the same time - when we define an operator which writes its output into a blob with a given name;
#  - NetBuilder keeps track of all operator outputs seen before current execution point in the same block and up the stack in parent blocks.
#  - If an operator writes into a previously unseen blob, it creates a **local** blob that is visible only within the current block and the subsequent children blocks. Local blobs created in a given block are effectively deleted when we exit the block. Any write into previously defined (in the same block or in the parent blocks) blob updates an originally created blob and does not result in the redefinition of a blob.
#  - Operator's input blobs have to be defined earlier in the same block or in the parent blocks up the stack. 

# As a result, in order to see values computed by a block after the block is finished the corresponding blobs have to be defined outside of the block. Note, that this is one of the ways to solve the problem with uninitialized blobs (e.g. blob created by 'then' branch, but not by 'else' branch), these rules effectively force visible blobs to be always correctly initialized.

# To illustrate concepts of block semantics and provide a more sophisticated example, let's consider the following net:

# In[4]:


with NetBuilder() as nb:
    ops.Const(0.0, blob_out="zero")
    ops.Const(1.0, blob_out="one")
    ops.Const(2.0, blob_out="two")
    ops.Const(1.5, blob_out="x")
    ops.Const(0.0, blob_out="y")
    with ops.IfNet(ops.GT(["x", "zero"])):
        ops.Copy("x", "local_blob")
        with ops.IfNet(ops.LE(["local_blob", "one"])):
            ops.Copy("one", "y")
        with ops.Else():
            ops.Copy("two", "y")
    with ops.Else():
        ops.Copy("zero", "y")
        # ops.Copy("local_blob", "z") - fails with exception during net construction, local_blob is undefined


# In[5]:


plan = Plan('if_net_test_2')
plan.AddStep(to_execution_step(nb))
ws = workspace.C.Workspace()
ws.run(plan)
print('x = ', ws.blobs["x"].fetch())
print('y = ', ws.blobs["y"].fetch())
assert "local_blob" not in ws.blobs


# Brew is another Caffe2 interface used to construct nets. Unlike NetBuilder, Brew does not track the hierarchy of blocks and, as a result, we need to specify which blobs are considered local and which global when passing 'then' and 'else' models to an API call:

# In[6]:


from caffe2.python import brew
from caffe2.python.workspace import FeedBlob, RunNetOnce, FetchBlob
from caffe2.python.model_helper import ModelHelper


# In[7]:


model = ModelHelper(name="test_if_model")

model.param_init_net.ConstantFill([], ["zero"], shape=[1], value=0.0)
model.param_init_net.ConstantFill([], ["one"], shape=[1], value=1.0)
model.param_init_net.ConstantFill([], ["x"], shape=[1], value=0.5)
model.param_init_net.ConstantFill([], ["y"], shape=[1], value=0.0)
model.param_init_net.GT(["x", "zero"], "cond")

then_model = ModelHelper(name="then_test_model")
then_model.net.Copy("one", "y")

else_model = ModelHelper(name="else_test_model")
else_model.net.Copy("zero", "y")

brew.cond(
    model=model,
    cond_blob="cond", # blob with condition value
    external_blobs=["x", "y", "zero", "one"], # writes into these blobs update existing blobs
    then_model=then_model,
    else_model=else_model)


# To run resulting init and main net:

# In[8]:


RunNetOnce(model.param_init_net)
RunNetOnce(model.net)
print("x = ", FetchBlob("x"))
print("y = ", FetchBlob("y"))


# ### Loops

# Another important control flow operator is 'While' that allows repeated execution of a fragment of net. Let's consider NetBuilder's version of While first.

# In[9]:


with NetBuilder() as nb:
    ops.Const(0, blob_out="i")
    ops.Const(0, blob_out="y")
    with ops.WhileNet():
        with ops.Condition():
            ops.Add(["i", ops.Const(1)], ["i"])
            ops.LE(["i", ops.Const(7)])
        ops.Add(["i", "y"], ["y"])


# This example illustrates the usage of 'while' loop with NetBuilder. As with an 'If' operator, standard block semantic rules apply. Note the usage of ops.Condition clause that should immediately follow ops.WhileNet and contains code that is executed before each iteration. The last operator in the condition clause is expected to have a single boolean output that determines whether the other iteration is executed.

# In the example above we increment the counter ("i") before each iteration and accumulate its values in "y" blob, the loop's body is executed 7 times, the resulting blob values:

# In[10]:


plan = Plan('while_net_test')
plan.AddStep(to_execution_step(nb))
ws = workspace.C.Workspace()
ws.run(plan)
print("i = ", ws.blobs["i"].fetch())
print("y = ", ws.blobs["y"].fetch())


# Corresponding Brew example:

# In[11]:


model = ModelHelper(name="test_while_model")

model.param_init_net.ConstantFill([], ["i"], shape=[1], value=0)
model.param_init_net.ConstantFill([], ["one"], shape=[1], value=1)
model.param_init_net.ConstantFill([], ["seven"], shape=[1], value=7)
model.param_init_net.ConstantFill([], ["y"], shape=[1], value=0)

loop_model = ModelHelper(name="loop_test_model")
loop_model.net.Add(["i", "y"], ["y"])

cond_model = ModelHelper(name="cond_test_model")
cond_model.net.Add(["i", "one"], "i")
cond_model.net.LE(["i", "seven"], "cond")

brew.loop(
    model=model,
    cond_blob="cond", # explicitly specifying condition blob
    external_blobs=["cond", "i", "one", "seven", "y"],
    loop_model=loop_model,
    cond_model=cond_model # condition model is optional
)


# Corresponding blob values:

# In[12]:


RunNetOnce(model.param_init_net)
RunNetOnce(model.net)
print("i = ", FetchBlob("i"))
print("y = ", FetchBlob("y"))


# ### Backpropagation

# Both 'If' and 'While' operators support backpropagation. To illustrate how backpropagation with control ops work, let's consider the following example:

# In[14]:


import numpy as np
# _use_control_ops=True forces NetBuilder to output single net as a result
# x is external for NetBuilder, so letting nb know about it through initial_scope param
FeedBlob("x", np.array(0.5, dtype='float32'))
with NetBuilder(_use_control_ops=True, initial_scope=["x"]) as nb:
    ops.Const(0.0, blob_out="zero")
    ops.Const(1.0, blob_out="one")
    ops.Const(4.0, blob_out="y")
    ops.Const(0.0, blob_out="z")
    with ops.IfNet(ops.GT(["x", "zero"])):
        ops.Pow("y", "z", exponent=2.0)
    with ops.Else():
        ops.Pow("y", "z", exponent=3.0)

assert len(nb.get()) == 1, "Expected a single net produced"
net = nb.get()[0]
grad_map = net.AddGradientOperators(["z"])


# Output blob "z" as a function of "y" depends on the value of blob "x", if "x" is greater than zero, than "z = y^2", otherwise it is "z = y^3"

# In[15]:


RunNetOnce(net)
print("x = ", FetchBlob("x"))
print("y = ", FetchBlob("y"))
print("z = ", FetchBlob("z"))
print("y_grad = ", FetchBlob("y_grad"))


# Now, let's change value of blob "x" and rerun net:

# In[16]:


FeedBlob("x", np.array(-0.5, dtype='float32'))
RunNetOnce(net)
print("x = ", FetchBlob("x"))
print("y = ", FetchBlob("y"))
print("z = ", FetchBlob("z"))
print("y_grad = ", FetchBlob("y_grad"))


# An example illustrating backpropagation on the loop:

# In[17]:


with NetBuilder(_use_control_ops=True) as nb:
    ops.Copy(ops.Const(0), "i")
    ops.Copy(ops.Const(1), "one")
    ops.Copy(ops.Const(2), "two")
    ops.Copy(ops.Const(2.0), "x")
    ops.Copy(ops.Const(3.0), "y")
    ops.Copy(ops.Const(2.0), "z")
    # computes x^4, y^2, z^3
    with ops.WhileNet():
        with ops.Condition():
            ops.Add(["i", "one"], "i")
            ops.LE(["i", "two"])
        ops.Pow("x", "x", exponent=2.0)
        with ops.IfNet(ops.LT(["i", "two"])):
            ops.Pow("y", "y", exponent=2.0)
        with ops.Else():
            ops.Pow("z", "z", exponent=3.0)

    ops.Add(["x", "y"], "x_plus_y")
    ops.Add(["x_plus_y", "z"], "s")

assert len(nb.get()) == 1, "Expected a single net produced"
net = nb.get()[0]

grad_map = net.AddGradientOperators(["s"])


# In[18]:


workspace.RunNetOnce(net)
print("x = ", FetchBlob("x"))
print("x_grad = ", FetchBlob("x_grad")) # 4x^3
print("y = ", FetchBlob("y"))
print("y_grad = ", FetchBlob("y_grad")) # 2y
print("z = ", FetchBlob("z"))
print("z_grad = ", FetchBlob("z_grad")) # 3z^2


# ### Implementation Notes

# On the low level, Caffe2 uses the following set of operators to implement forward and backward branching and loops:
# - If - accepts *then_net* and *else_net* nets as arguments and executes one of them, depending on input condition blob value, nets are executed **in the same** workspace;
# - While - repeats execution of *loop_net* net passed as argument, net is executed in the same workspace;
# - Do - special operator that creates a separate inner workspace, setups blob mappings between outer and inner workspaces and runs a net in an inner workspace;
# - CreateScope/HasScope - special operators that create and keep track of workspaces used by Do operator.
# 
# Higher level libraries that implement branching and looping (e.g. in NetBuilder, Brew), use these operators to build control flow, e.g. for 'If':
#  - do necessary sanity checks (e.g. determine which blobs are initialized and check that subnet does not read undefined blobs)
#  - wrap 'then' and 'else' branches into Do
#  - setup correct blob mappings by specifying which local names are mapped to outer blobs
#  - prepare scope structure, used by Do operator
# 
# While 'If' and 'While' Caffe2 ops can be used directly without creating local block workspaces, we encourage users to use higher level Caffe2 interfaces that provide necessary correctness guarantees.
# 
# Backpropagation for 'While' in general is expensive memory-wise - we have to save local workspace for every iteration of a block, including global blobs visible to the block. It is recommended that users use RecurrentNetwork operator instead in production environments.

