| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import schema |
| | from caffe2.python.layers.layers import ( |
| | ModelLayer, |
| | ) |
| |
|
| |
|
| | class Split(ModelLayer): |
| |
|
| | def __init__(self, model, input_record, num_splits=1, axis=1, |
| | name='split', split=None, **kwargs): |
| | super(Split, self).__init__(model, name, input_record, **kwargs) |
| | self.axis = axis |
| | |
| | |
| | axis -= 1 |
| | assert axis >= 0 |
| |
|
| | assert isinstance(input_record, schema.Scalar),\ |
| | "Incorrect input type. Expected Scalar, but received: {0}".\ |
| | format(input_record) |
| |
|
| | input_shape = input_record.field_type().shape |
| | assert len(input_shape) >= axis |
| | if split is None: |
| | assert input_shape[axis] % num_splits == 0 |
| | else: |
| | num_splits = len(split) |
| | assert input_shape[axis] == sum(split) |
| |
|
| | if split is None: |
| | output_shape = list(input_shape) |
| | output_shape[axis] = int(output_shape[axis] / num_splits) |
| | else: |
| | output_shape = [] |
| | for i in range(num_splits): |
| | output_shape_i = list(input_shape) |
| | output_shape_i[axis] = split[i] |
| | output_shape.append(output_shape_i) |
| |
|
| | data_type = input_record.field_type().base |
| |
|
| |
|
| | if split is None: |
| | output_scalars = [ |
| | schema.Scalar( |
| | (data_type, output_shape), |
| | self.get_next_blob_reference('output_{}'.format(i)), |
| | ) |
| | for i in range(num_splits) |
| | ] |
| | else: |
| | output_scalars = [ |
| | schema.Scalar( |
| | (data_type, output_shape[i]), |
| | self.get_next_blob_reference('output_{}'.format(i)), |
| | ) |
| | for i in range(num_splits) |
| | ] |
| | self.output_schema = schema.Tuple(*output_scalars) |
| | self.split = split |
| |
|
| | def add_ops(self, net): |
| | net.Split( |
| | self.input_record.field_blobs(), |
| | self.output_schema.field_blobs(), |
| | split=self.split, |
| | axis=self.axis, |
| | ) |
| |
|
