| from caffe2.python import workspace, core |
| import numpy as np |
| |
| from utils import NUM_LOOP_ITERS |
| |
| workspace.GlobalInit(['caffe2']) |
| |
| def add_blob(ws, blob_name, tensor_size): |
| blob_tensor = np.random.randn(*tensor_size).astype(np.float32) |
| ws.FeedBlob(blob_name, blob_tensor) |
| |
| class C2SimpleNet(object): |
| """ |
| This module constructs a net with 'op_name' operator. The net consist |
| a series of such operator. |
| It initializes the workspace with input blob equal to the number of parameters |
| needed for the op. |
| Provides forward method to run the net niter times. |
| """ |
| def __init__(self, op_name, num_inputs=1, debug=False): |
| self.input_names = [] |
| self.net = core.Net("framework_benchmark_net") |
| self.input_names = ["in_{}".format(i) for i in range(num_inputs)] |
| for i in range(num_inputs): |
| add_blob(workspace, self.input_names[i], [1]) |
| self.net.AddExternalInputs(self.input_names) |
| op_constructor = getattr(self.net, op_name) |
| op_constructor(self.input_names) |
| self.output_name = self.net._net.op[-1].output |
| print("Benchmarking op {}:".format(op_name)) |
| for _ in range(NUM_LOOP_ITERS): |
| output_name = self.net._net.op[-1].output |
| self.input_names[-1] = output_name[0] |
| assert len(self.input_names) == num_inputs |
| op_constructor(self.input_names) |
| workspace.CreateNet(self.net) |
| if debug: |
| print(self.net._net) |
| |
| def forward(self, niters): |
| workspace.RunNet(self.net, niters, False) |
