/usr/local/lib64/python3.6/site-packages/caffe2/python/layers
Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/layers/layers.py (17412B)
## @package layers
# Module caffe2.python.layers.layers
import logging
from collections import namedtuple
import numpy as np
from caffe2.proto import caffe2_pb2
from caffe2.python import core, schema, scope, utils, workspace
from caffe2.python.layers.tags import TagContext
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# Some types to simplify descriptions of things traveling between ops
IdList = schema.List(np.int64)
IdScoreList = schema.Map(np.int64, np.float32)
IdListWithEvicted = schema.ListWithEvicted(np.int64)
IdScoreListWithEvicted = schema.MapWithEvicted(np.int64, np.float32)
def almost_equal_schemas(
record,
original_schema,
check_field_names=True,
check_field_types=True,
check_field_metas=False,
):
if original_schema == IdList:
return schema.equal_schemas(
record,
IdList,
check_field_names=check_field_names,
check_field_types=check_field_types,
check_field_metas=check_field_metas,
) or schema.equal_schemas(
record,
IdListWithEvicted,
check_field_names=check_field_names,
check_field_types=check_field_types,
check_field_metas=check_field_metas,
)
elif original_schema == IdScoreList:
return schema.equal_schemas(
record,
IdScoreList,
check_field_names=check_field_names,
check_field_types=check_field_types,
check_field_metas=check_field_metas,
) or schema.equal_schemas(
record,
IdScoreListWithEvicted,
check_field_names=check_field_names,
check_field_types=check_field_types,
check_field_metas=check_field_metas,
)
else:
return schema.equal_schemas(record, original_schema)
def get_key(record):
if almost_equal_schemas(record, IdList):
key = "values"
elif almost_equal_schemas(
record, IdScoreList, check_field_types=False
):
key = "values:keys"
else:
raise NotImplementedError("Not implemented for {}".format(record))
assert record[key].metadata is not None, "Blob {} doesn't have metadata".format(
str(record[key]())
)
return record[key]
def get_categorical_limit(record):
key = get_key(record)
return key.metadata.categorical_limit
def get_avg_length(record):
return record["lengths"].metadata.expected_value
def set_request_only(field):
for f in field.all_scalars():
categorical_limit, expected_value = None, None
if not f.metadata:
feature_specs = schema.FeatureSpec(feature_is_request_only=True)
elif not f.metadata.feature_specs:
categorical_limit = f.metadata.categorical_limit
expected_value = f.metadata.expected_value
feature_specs = schema.FeatureSpec(feature_is_request_only=True)
else:
categorical_limit = f.metadata.categorical_limit
expected_value = f.metadata.expected_value
feature_specs = schema.FeatureSpec(
feature_type=f.metadata.feature_specs.feature_type,
feature_names=f.metadata.feature_specs.feature_names,
feature_ids=f.metadata.feature_specs.feature_ids,
feature_is_request_only=True,
desired_hash_size=f.metadata.feature_specs.desired_hash_size,
)
# make sure not to set categorical_limit for a non-integer field
if not np.issubdtype(f.field_type(), np.integer):
assert (
categorical_limit is None
), "categorical_limit shouldn't be set for no-integer field"
f.set_metadata(
schema.Metadata(
categorical_limit=categorical_limit,
expected_value=expected_value,
feature_specs=feature_specs,
)
)
class InstantiationContext(object):
"""
List of contexts where layer could be instantitated
"""
# The layers support this context will accumulate predictions, labels,
# weights. The accumulated data can later be used to compute
# calibration or for other
# purpose.
ACCUMULATE_PRED = "accumulate_pred"
EVAL = "eval"
PREDICTION = "prediction"
TRAINING = "training"
_LAYER_REGISTRY = {}
def register_layer(name, layer):
assert name not in _LAYER_REGISTRY, "{0} already exists".format(name)
_LAYER_REGISTRY[name] = layer
def layer_exists(name):
return name in _LAYER_REGISTRY
def get_layer_class(name):
return _LAYER_REGISTRY[name]
def create_layer(layer_name, *args, **kwargs):
return _LAYER_REGISTRY[layer_name](*args, **kwargs)
LayerPsParam = namedtuple("LayerPsParam", ["sparse_key", "average_length"])
class LayerParameter(object):
def __init__(
self,
parameter=None,
optimizer=None,
initializer=None,
ps_param=None,
regularizer=None,
):
assert isinstance(
parameter, core.BlobReference
), "expect {0} to be a blob reference".format(str(parameter))
# need to put the following line (shape) before initialier
# shape will be updated once initializer is (re)set
self._shape = None
self.parameter = parameter
self.optimizer = optimizer
self.initializer = initializer
self.ps_param = ps_param
self.regularizer = regularizer
@property
def initializer(self):
return self._initializer
@initializer.setter
def initializer(self, op):
assert op is None or core.IsOperator(
getattr(op, "type", None)
), "initializer expects an operator, got type: {}".format(type(op))
self._initializer = op
if op is not None:
self.shape = self._infer_shape_from_initializer()
@property
def shape(self):
return self._shape
@shape.setter
def shape(self, shape):
assert self.shape is None or self.shape == shape, (
"inconsistent shape for layer parameter:"
" {}, expect: {}, but got {}".format(self, self.shape, shape)
)
self._shape = shape
def _infer_shape_from_initializer(self):
for arg in self.initializer.arg:
if arg.name == "shape":
return list(arg.ints)
with workspace.WorkspaceGuard("model_init_by_loading_params"):
try:
net = core.Net("shape_checker")
net._net.op.extend([self.initializer])
shape_blob = net.NextScopedBlob(self.parameter + "_shape")
net.Shape([self.parameter], shape_blob)
workspace.RunNetOnce(net)
shape = workspace.FetchBlob(shape_blob).tolist()
# ResetWorkspace to save memory
workspace.ResetWorkspace()
return shape
except RuntimeError as exp:
logger.warning(
"Cannot infer the shape of blob {} from operator {}: {}".format(
self.parameter, self.initializer.type, exp
)
)
workspace.ResetWorkspace()
return None
def __str__(self):
return str(self.parameter)
def is_request_only_scalar(scalar):
if len(scalar.field_metadata()) == 0:
return False
for metadata in scalar.field_metadata():
if not (
metadata
and metadata.feature_specs
and getattr(metadata.feature_specs, "feature_is_request_only", False)
):
return False
return True
# Contains features accessed in a model layer of a given type
# `type`: A string representing the kind of feature, consistent with FeatureSpec
# `ids`: A set of feature IDs that are accessed in the model layer
AccessedFeatures = namedtuple("AccessedFeatures", ["type", "ids"])
class ModelLayer(object):
def __init__(
self,
model,
prefix,
input_record,
predict_input_record_fields=None,
tags=None,
**kwargs
):
"""
Base class for model layers. Layer is an abstraction that allows to
provide model description in terms of meta-operators, where each of the
meta-operators can have different implementations for training,
evaluation and prediction, that are instantiated later. As an example
SampledSoftmax can do something related to sampling depending on
supervision during the training and just apply softmax if it's used for
prediction/evaluation.
All inputs/outputs from layers are represented as a record (instance of
schema bounded to blobs) and are accessible through input_record and
output_schema. If Layer needs to have only a subset of inputs/provides
subset of outputs during the inference - it should provide
predict_input_record and predict_output_schema correspondingly (those
records are expected to be a subset of input_record/output_schema).
Each layer has a list of Tags associated with it, that depends on
current context and arguments. It's possible to use those tags during
the instantiation time.
"""
self.name = model.next_layer_name(prefix)
self.model = model
self.kwargs = kwargs
self._input_record = input_record
if predict_input_record_fields:
if not isinstance(predict_input_record_fields, list):
predict_input_record_fields = [predict_input_record_fields]
self._predict_input_record = self._input_record[predict_input_record_fields]
else:
self._predict_input_record = None
self.request_only = True
if len(input_record.all_scalars()) == 0:
self.request_only = False
for scalar in input_record.all_scalars():
if not is_request_only_scalar(scalar):
self.request_only = False
break
self.precomputation_request_only = False
self.precomputation_object_only = False
self._output_schema = None
self._predict_output_schema = None
self.eval_output_schema = None
self.tags = set(tags or [])
self.tags.update(TagContext.current().tags)
self.params = []
self._export_output_for_metrics = False
self._export_params_for_metrics = False
def get_type(self):
return self.__class__.__name__
def _check_output_schema(self):
assert self._output_schema is not None, "Schema is not initialized"
assert self._predict_output_schema is None or schema.is_schema_subset(
self._predict_output_schema, self._output_schema
), "predict_output_schema is not a subset of the output_schema"
@property
def predict_input_record(self):
return self._predict_input_record or self._input_record
@property
def input_record(self):
return self._input_record
@property
def predict_output_schema(self):
self._check_output_schema()
return self._predict_output_schema or self._output_schema
@predict_output_schema.setter
def predict_output_schema(self, output_schema):
assert self._predict_output_schema is None
self._predict_output_schema = output_schema
@property
def output_schema(self):
if self.request_only:
set_request_only(self._output_schema)
self._check_output_schema()
return self._output_schema
@output_schema.setter
def output_schema(self, output_schema):
assert self._output_schema is None
self._output_schema = output_schema
def get_parameters(self):
return self.params
def get_fp16_compatible_parameters(self):
"""Return a subset of parameters which can be converted to fp16"""
return []
def get_memory_usage(self):
return 0
def get_accessed_features(self):
"""
Return a map from field to list of AccessedFeatures, the map should
contain all features accessed in the model layer
"""
return {}
def add_init_params(self, init_net):
"""
Adds layer initialization operators to passed net.
"""
for param in self.params:
# TODO(amalevich): Either return back to lambdas, that add
# all params (looks a bit safer and breaking less
# abstractions) or extend Net interface to this type of
# operations better
# TODO(xlwang) init_net._net.op has type google.protobuf.\
# internal.containers.RepeatedCompositeFieldContainer, but
# the version of protobuf in fbcode does not support append
# so extend is used
init_op = param.initializer
current_device_scope = scope.CurrentDeviceScope()
if not init_op:
continue
if not init_op.HasField("device_option") and current_device_scope:
init_op = caffe2_pb2.OperatorDef()
init_op.CopyFrom(param.initializer)
init_op.device_option.CopyFrom(current_device_scope)
# do not add duplicated init ops
if any(
utils.OpAlmostEqual(op, init_op, "debug_info")
for op in init_net._net.op
):
continue
init_net._net.op.extend([init_op])
def create_param(
self, param_name, shape, initializer, optimizer, ps_param=None, regularizer=None
):
with scope.NameScope(self.name, reset=True):
param = self.model.create_param(
param_name=param_name,
shape=shape,
initializer=initializer,
optimizer=optimizer,
ps_param=ps_param,
regularizer=regularizer,
)
# make sure we don't share parameters in the same layer
assert all(param.parameter != p.parameter for p in self.params)
self.params.append(param)
return param.parameter
def get_next_blob_reference(self, name):
with scope.NameScope(self.name, reset=True):
return self.model.net.NextScopedBlob(name)
def add_operators(self, net, init_net=None, context=InstantiationContext.TRAINING):
"""
Adds layer trainig or initialization operators to the passed in net.
init_net can be None and can be called independently from add_init_params
"""
# Namescope below should warranty that all intermediate blobs will be
# assiciated with the layer that produces them
with scope.NameScope(self.name):
if context not in {
InstantiationContext.PREDICTION,
InstantiationContext.EVAL,
InstantiationContext.ACCUMULATE_PRED,
}:
assert init_net, "Only prediction and eval context don't need init_net"
if init_net:
self.add_init_params(init_net)
if context == InstantiationContext.TRAINING:
self.add_train_ops(net)
elif context == InstantiationContext.EVAL:
self.add_eval_ops(net)
elif context == InstantiationContext.ACCUMULATE_PRED:
self.add_ops_to_accumulate_pred(net)
else:
self.add_ops(net)
if (
context in {InstantiationContext.TRAINING, InstantiationContext.EVAL}
and self._export_params_for_metrics
):
self.add_param_copy_operators(net)
def add_ops(self, net):
# Predict layer implementation.
raise NotImplementedError
def add_eval_ops(self, net):
# Default eval layer implementation is completely matching
# predict layer implementation.
self.add_ops(net)
def add_train_ops(self, net):
# Default train layer implementation is completely matching
# eval layer implementation.
self.add_eval_ops(net)
def add_ops_to_accumulate_pred(self, net):
# This adds operators to accumulate predictions/labels/weights. The
# accumulated data can later be used to compute calibration or for other
# purpose. Default layer implementation is completely matching eval
# layer implementation.
self.add_eval_ops(net)
def add_param_copy_operators(self, net):
for param in self.params:
param_copy_ref = self.model.metrics_schema[str(param.parameter)]
net.Copy([param.parameter], param_copy_ref.field_blobs())
def export_output_for_metrics(self):
self._export_output_for_metrics = True
# Export output of the layer directly
export_name = self.name + "/output"
self.model.add_metric_field(export_name, self.output_schema)
def export_params_for_metrics(self):
self._export_params_for_metrics = True
# Export copies of parameters
for param in self.params:
param_copy_ref = self.get_next_blob_reference(
str(param).split("/")[-1] + "_copy"
)
self.model.add_metric_field(str(param.parameter), param_copy_ref)