| from caffe2.python import core, workspace | |
| from caffe2.python.test_util import TestCase | |
| import numpy as np | |
| import unittest | |
| class DoOpTest(TestCase): | |
| def test_operator(self): | |
| def make_net(): | |
| subnet = core.Net('subnet') | |
| subnet.Add(["X", "Y"], "Z") | |
| net = core.Net("net") | |
| net.CreateScope([], "W") | |
| net.Do( | |
| ["outer_X", "outer_Y", "W"], | |
| ["outer_Z", "W"], | |
| net=subnet.Proto(), | |
| inner_blobs=["X", "Y", "Z"], | |
| outer_blobs_idx=[0, 1, 2], | |
| ) | |
| return net | |
| net = make_net() | |
| workspace.ResetWorkspace() | |
| workspace.FeedBlob("outer_X", np.asarray([1, 2])) | |
| workspace.FeedBlob("outer_Y", np.asarray([3, 4])) | |
| workspace.RunNetOnce(net) | |
| outer_Z_val = workspace.FetchBlob("outer_Z") | |
| self.assertTrue(np.all(outer_Z_val == np.asarray([4, 6]))) | |
| def test_reuse_workspace(self): | |
| def make_net(): | |
| param_init_subnet = core.Net('param_init_subnet') | |
| param_init_subnet.ConstantFill([], "X", shape=[1], value=1) | |
| param_init_subnet.ConstantFill([], "Y", shape=[1], value=2) | |
| subnet = core.Net("subnet") | |
| subnet.Add(["X", "Y"], "Z") | |
| net = core.Net("net") | |
| net.CreateScope([], "W") | |
| net.Do( | |
| "W", "W", | |
| net=param_init_subnet.Proto(), | |
| inner_blobs=[], | |
| outer_blobs_idx=[], | |
| ) | |
| net.Do( | |
| "W", ["outer_Z", "W"], | |
| net=subnet.Proto(), | |
| inner_blobs=["Z"], | |
| outer_blobs_idx=[0], | |
| reuse_workspace=True, | |
| ) | |
| return net | |
| net = make_net() | |
| workspace.ResetWorkspace() | |
| workspace.RunNetOnce(net) | |
| outer_Z_val = workspace.FetchBlob("outer_Z") | |
| self.assertTrue(np.all(outer_Z_val == np.asarray([3]))) | |
| if __name__ == '__main__': | |
| unittest.main() | |