mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-06 00:54:56 +08:00
8286ce1e3a
Summary: Closes https://github.com/caffe2/caffe2/pull/1260 Differential Revision: D5906739 Pulled By: Yangqing fbshipit-source-id: e482ba9ba60b5337d9165f28f7ec68d4518a0902
119 lines
4.6 KiB
Python
119 lines
4.6 KiB
Python
# Copyright (c) 2016-present, Facebook, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
##############################################################################
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
from caffe2.python import core, 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.C.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 explicity."
|
|
)
|
|
output_names = get_name_list(
|
|
output_prefix, max_output, max_output
|
|
)
|
|
for i, input_blob in enumerate(inputs):
|
|
ws.create_blob(input_names[i]).feed(input_blob)
|
|
|
|
op = core.CreateOperator(
|
|
op_type, input_names, output_names, **args
|
|
)
|
|
ws._run_operator(op.SerializeToString())
|
|
# RunOperator
|
|
output_values = [ws.fetch_blob(x) for x in output_names]
|
|
return namedtupledict('output', output_names)(*output_values)
|
|
|
|
return op_func
|
|
|
|
|
|
Functional = _Functional()
|