| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import schema |
| | from caffe2.python.layers.layers import ( |
| | ModelLayer, |
| | ) |
| | from future.utils import viewitems |
| | import numpy as np |
| | from collections import defaultdict |
| |
|
| | import logging |
| | logger = logging.getLogger(__name__) |
| |
|
| |
|
| | def get_concatenated_feature_to_index(blobs_to_concat): |
| | concat_feature_to_index = defaultdict(list) |
| | start_pos = 0 |
| | for scalar in blobs_to_concat: |
| | num_dims = scalar.dtype.shape[0] |
| | if hasattr(scalar, 'metadata') \ |
| | and hasattr(scalar.metadata, 'feature_specs') \ |
| | and hasattr(scalar.metadata.feature_specs, 'feature_to_index') \ |
| | and isinstance(scalar.metadata.feature_specs.feature_to_index, dict): |
| | for k, v in scalar.metadata.feature_specs.feature_to_index.items(): |
| | concat_feature_to_index[k].extend([start_pos + vi for vi in v]) |
| | start_pos += num_dims |
| | return dict(concat_feature_to_index) if concat_feature_to_index.keys() else None |
| |
|
| |
|
| | class Concat(ModelLayer): |
| | """ |
| | Construct Concat layer |
| | Assume that first dimension is batch, |
| | |
| | Example: |
| | |
| | embedding_dim = 64 |
| | input_record = self.new_record(schema.Struct( |
| | ('input1', schema.Scalar((np.float32, (embedding_dim, )))), |
| | ('input2', schema.Scalar((np.float32, (embedding_dim, )))), |
| | ('input3', schema.Scalar((np.float32, (embedding_dim, )))), |
| | )) |
| | |
| | output = self.model.Concat(input_record) |
| | self.assertEqual( |
| | schema.Scalar((np.float32, ((len(input_record.fields) * embedding_dim, )))), |
| | output |
| | ) |
| | |
| | # Note that in Concat layer we assume first dimension is batch. |
| | # so input is B * embedding_dim |
| | # add_axis=1 make it B * 1 * embedding_dim |
| | # Concat on axis=1 make it B * N * embedding_dim |
| | |
| | output = self.model.Concat(input_record, axis=1, add_axis=1) |
| | self.assertEqual( |
| | schema.Scalar((np.float32, ((len(input_record.fields), embedding_dim)))), |
| | output |
| | ) |
| | """ |
| |
|
| | def __init__(self, model, input_record, axis=1, add_axis=0, |
| | name='concat', **kwargs): |
| | super(Concat, self).__init__(model, name, input_record, **kwargs) |
| | self.axis = axis |
| | self.add_axis = add_axis |
| | assert not (axis == 0 and add_axis == 1), \ |
| | "It's not allowed to add axis=0" |
| | assert isinstance(input_record, schema.Struct),\ |
| | "Incorrect input type. Expected Struct, but received: {0}".\ |
| | format(input_record) |
| |
|
| | shapes = [] |
| | for field_name, field_type in viewitems(input_record.fields): |
| | assert isinstance(field_type, schema.Scalar),\ |
| | "Incorrect input type for {}. Expected Scalar, but got: {}".\ |
| | format(field_name, field_type) |
| | |
| | |
| | shape = list(field_type.field_type().shape) |
| | if add_axis: |
| | shape.insert(axis - 1, 1) |
| | assert len(shape) >= axis,\ |
| | "Concat expects that limited dimensions of the input tensor" |
| | shapes.append(shape) |
| | logger.info('Concat Layer input shapes: ' + str(shapes)) |
| |
|
| | if axis == 0: |
| | self.output_schema = schema.from_blob_list( |
| | input_record[0], |
| | [self.get_next_blob_reference('output')] |
| | ) |
| | return |
| |
|
| | concat_dim = 0 |
| | for shape in shapes: |
| | concat_dim += shape[axis - 1] |
| | shape[axis - 1] = 0 |
| | assert shape == shapes[0],\ |
| | "Shapes {0} and {1} are not compatible for Concat".\ |
| | format(shape, shapes[0]) |
| | output_dims = shapes[0] |
| | output_dims[axis - 1] = concat_dim |
| |
|
| | logger.info('Concat Layer output_dims: ' + str(output_dims)) |
| | self.output_schema = schema.Scalar( |
| | (np.float32, output_dims), |
| | self.get_next_blob_reference('output')) |
| |
|
| | record_to_concat = input_record.fields.values() |
| | concated_feature_to_index = get_concatenated_feature_to_index( |
| | record_to_concat |
| | ) |
| | if concated_feature_to_index: |
| | metadata = schema.Metadata( |
| | feature_specs=schema.FeatureSpec( |
| | feature_to_index=concated_feature_to_index |
| | ) |
| | ) |
| | self.output_schema.set_metadata(metadata) |
| |
|
| |
|
| | def add_ops(self, net): |
| | net.Concat( |
| | self.input_record.field_blobs(), |
| | [ |
| | self.output_schema.field_blobs()[0], |
| | self.output_schema.field_blobs()[0] + "_concat_dims" |
| | ], |
| | axis=self.axis, |
| | add_axis=self.add_axis, |
| | ) |
| |
|
