1 from __future__
import absolute_import
2 from __future__
import division
3 from __future__
import print_function
4 from __future__
import unicode_literals
7 from caffe2.proto
import caffe2_pb2
11 def rewrite_init_net_simple(net):
13 op.device_option.device_type = caffe2_pb2.MKLDNN
15 def last_producer(ops, blob):
16 for (i, op)
in reversed(list(enumerate(ops))):
19 raise ValueError(
"Failed to find last producer of blob, %s", blob)
22 def rewrite_run_net_simple(net):
27 return "{}__MKL__".format(name)
29 input_blob = net.external_input[0]
30 if input_blob != net.op[0].input[0]:
32 "Input blob: {} is not consumed by first op: {}".format(
33 input_blob, net.op[0]))
35 copy_input_op = core.CreateOperator(
36 "CopyCPUToMKL", input_blob, mkl_tmp(input_blob))
37 net.op[0].input[0] = mkl_tmp(input_blob)
40 core.CreateOperator(
"CopyMKLToCPU", mkl_tmp(output_blob), output_blob)
41 for output_blob
in net.external_output]
43 for output_blob
in net.external_output:
44 last_producer_idx = last_producer(net.op, output_blob)
45 renamed_outputs = [blob
if blob != output_blob
else mkl_tmp(blob)
46 for blob
in net.op[last_producer_idx].output]
47 net.op[last_producer_idx].output[:] = renamed_outputs
49 for op
in net.op[last_producer_idx + 1:]:
50 renamed_input = [blob
if blob != output_blob
else mkl_tmp(blob)
52 op.input[:] = renamed_input
54 ops = [copy_input_op] + net.op[:] + copy_output_ops
58 op.device_option.MergeFrom(
59 core.DeviceOption(device_type=caffe2_pb2.MKLDNN))
63 def rewrite_model_helper_simple(model):
64 model = copy.deepcopy(model)
66 rewrite_init_net_simple(model.param_init_net.Proto())
67 rewrite_run_net_simple(model.net.Proto())