115 lines
4.3 KiB
Python
115 lines
4.3 KiB
Python
|
|
|
|
|
|
|
|
|
|
from caffe2.python import core, workspace
|
|
from caffe2.proto import caffe2_pb2
|
|
from caffe2.python.onnx.workspace import Workspace
|
|
from collections import namedtuple
|
|
from six import string_types
|
|
|
|
OpSchema = workspace.C.OpSchema
|
|
|
|
|
|
def namedtupledict(typename, field_names, *args, **kwargs):
|
|
field_names_map = {n: i for i, n in enumerate(field_names)}
|
|
# Some output names are invalid python identifier, e.g. "0"
|
|
kwargs.setdefault('rename', True)
|
|
data = namedtuple(typename, field_names, *args, **kwargs)
|
|
|
|
def getitem(self, key):
|
|
if isinstance(key, string_types):
|
|
key = field_names_map[key]
|
|
return super(type(self), self).__getitem__(key)
|
|
|
|
data.__getitem__ = getitem
|
|
return data
|
|
|
|
|
|
class _Functional(object):
|
|
def __getattribute__(self, op_type):
|
|
def op_func(*inputs, **args):
|
|
ws = Workspace()
|
|
schema = OpSchema.get(op_type)
|
|
input_prefix = 'input_'
|
|
output_prefix = 'output_'
|
|
|
|
def get_name_list(prefix, num, max_num):
|
|
return [prefix + str(x) for x in range(min(num, max_num))]
|
|
|
|
input_names, output_names = [], []
|
|
input_names = get_name_list(
|
|
input_prefix, len(inputs), schema.max_input
|
|
)
|
|
# verify the length of input name is in range
|
|
# of schema
|
|
num_input = len(input_names)
|
|
if num_input > schema.max_input or num_input < \
|
|
schema.min_input or not schema.num_inputs_allowed(num_input):
|
|
raise ValueError(
|
|
"Functional C2: Number of inputs not in \
|
|
range: {} - {} or not allowed."
|
|
.format(schema.min_input, schema.max_input)
|
|
)
|
|
|
|
if 'num_output' in args:
|
|
num_output = args['num_output']
|
|
if num_output > schema.max_output or \
|
|
num_output < schema.min_output or \
|
|
not schema.num_outputs_allowed(num_output) or \
|
|
not schema.num_inputs_outputs_allowed(num_input,
|
|
num_output):
|
|
raise ValueError(
|
|
"Functional C2: Number of output \
|
|
not in range: {} - {} or not allowed"
|
|
.format(schema.min_output, schema.max_output)
|
|
)
|
|
output_names = get_name_list(
|
|
output_prefix, num_output, schema.max_output
|
|
)
|
|
args.pop('num_output')
|
|
calculated = schema.CalculateOutput(num_input)
|
|
if not output_names and calculated != -1:
|
|
output_names = get_name_list(
|
|
output_prefix, calculated, schema.max_output
|
|
)
|
|
|
|
if not output_names:
|
|
max_output = schema.max_output
|
|
# For an op with max_output == inf
|
|
# and no Output defined in schema
|
|
# user should pass output_size explicitly
|
|
if schema.inf == max_output:
|
|
raise ValueError(
|
|
"For operators with max_output == inf,\
|
|
user should pass num_output explicitly."
|
|
)
|
|
output_names = get_name_list(
|
|
output_prefix, max_output, max_output
|
|
)
|
|
|
|
# There could be input-output inplace enforcement; replace the
|
|
# output names with input ones if such enforcements exist
|
|
for i in range(len(input_names)):
|
|
for j in range(len(output_names)):
|
|
if schema.inplace_enforced(i, j):
|
|
output_names[j] = input_names[i]
|
|
|
|
op = core.CreateOperator(
|
|
op_type, input_names, output_names, **args
|
|
)
|
|
device_option = args.get('device_option', core.DeviceOption(caffe2_pb2.CPU))
|
|
with core.DeviceScope(device_option):
|
|
for i, input_blob in enumerate(inputs):
|
|
ws.FeedBlob(input_names[i], input_blob)
|
|
# RunOperator
|
|
ws.RunOperatorOnce(op)
|
|
output_values = [ws.FetchBlob(x) for x in output_names]
|
|
return namedtupledict('output', output_names)(*output_values)
|
|
|
|
return op_func
|
|
|
|
|
|
Functional = _Functional()
|