/usr/local/lib64/python3.6/site-packages/caffe2/python
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/observer_test.py (5316B)
import numpy as np
import unittest
from hypothesis import given, settings
import hypothesis.strategies as st
from caffe2.python import brew, core, model_helper, rnn_cell
import caffe2.python.workspace as ws
class TestObservers(unittest.TestCase):
def setUp(self):
core.GlobalInit(["python", "caffe2"])
ws.ResetWorkspace()
self.model = model_helper.ModelHelper()
brew.fc(self.model, "data", "y",
dim_in=4, dim_out=2,
weight_init=('ConstantFill', dict(value=1.0)),
bias_init=('ConstantFill', dict(value=0.0)),
axis=0)
ws.FeedBlob("data", np.zeros([4], dtype='float32'))
ws.RunNetOnce(self.model.param_init_net)
ws.CreateNet(self.model.net)
def testObserver(self):
ob = self.model.net.AddObserver("TimeObserver")
ws.RunNet(self.model.net)
print(ob.average_time())
num = self.model.net.NumObservers()
self.model.net.RemoveObserver(ob)
assert(self.model.net.NumObservers() + 1 == num)
@given(
num_layers=st.integers(1, 4),
forward_only=st.booleans()
)
@settings(deadline=1000)
def test_observer_rnn_executor(self, num_layers, forward_only):
'''
Test that the RNN executor produces same results as
the non-executor (i.e running step nets as sequence of simple nets).
'''
Tseq = [2, 3, 4]
batch_size = 10
input_dim = 3
hidden_dim = 3
run_cnt = [0] * len(Tseq)
avg_time = [0] * len(Tseq)
for j in range(len(Tseq)):
T = Tseq[j]
ws.ResetWorkspace()
ws.FeedBlob(
"seq_lengths",
np.array([T] * batch_size, dtype=np.int32)
)
ws.FeedBlob("target", np.random.rand(
T, batch_size, hidden_dim).astype(np.float32))
ws.FeedBlob("hidden_init", np.zeros(
[1, batch_size, hidden_dim], dtype=np.float32
))
ws.FeedBlob("cell_init", np.zeros(
[1, batch_size, hidden_dim], dtype=np.float32
))
model = model_helper.ModelHelper(name="lstm")
model.net.AddExternalInputs(["input"])
init_blobs = []
for i in range(num_layers):
hidden_init, cell_init = model.net.AddExternalInputs(
"hidden_init_{}".format(i),
"cell_init_{}".format(i)
)
init_blobs.extend([hidden_init, cell_init])
output, last_hidden, _, last_state = rnn_cell.LSTM(
model=model,
input_blob="input",
seq_lengths="seq_lengths",
initial_states=init_blobs,
dim_in=input_dim,
dim_out=[hidden_dim] * num_layers,
drop_states=True,
forward_only=forward_only,
return_last_layer_only=True,
)
loss = model.AveragedLoss(
model.SquaredL2Distance([output, "target"], "dist"),
"loss"
)
# Add gradient ops
if not forward_only:
model.AddGradientOperators([loss])
# init
for init_blob in init_blobs:
ws.FeedBlob(init_blob, np.zeros(
[1, batch_size, hidden_dim], dtype=np.float32
))
ws.RunNetOnce(model.param_init_net)
# Run with executor
self.enable_rnn_executor(model.net, 1, forward_only)
np.random.seed(10022015)
input_shape = [T, batch_size, input_dim]
ws.FeedBlob(
"input",
np.random.rand(*input_shape).astype(np.float32)
)
ws.FeedBlob(
"target",
np.random.rand(
T,
batch_size,
hidden_dim
).astype(np.float32)
)
ws.CreateNet(model.net, overwrite=True)
time_ob = model.net.AddObserver("TimeObserver")
run_cnt_ob = model.net.AddObserver("RunCountObserver")
ws.RunNet(model.net)
avg_time[j] = time_ob.average_time()
run_cnt[j] = int(''.join(x for x in run_cnt_ob.debug_info() if x.isdigit()))
model.net.RemoveObserver(time_ob)
model.net.RemoveObserver(run_cnt_ob)
print(avg_time)
print(run_cnt)
self.assertTrue(run_cnt[1] > run_cnt[0] and run_cnt[2] > run_cnt[1])
self.assertEqual(run_cnt[1] - run_cnt[0], run_cnt[2] - run_cnt[1])
def enable_rnn_executor(self, net, value, forward_only):
num_found = 0
for op in net.Proto().op:
if op.type.startswith("RecurrentNetwork"):
for arg in op.arg:
if arg.name == 'enable_rnn_executor':
arg.i = value
num_found += 1
# This sanity check is so that if someone changes the
# enable_rnn_executor parameter name, the test will
# start failing as this function will become defective.
self.assertEqual(1 if forward_only else 2, num_found)