(self, model, input_record, num_splits=1, axis=1,
name='split', split=None, **kwargs)
| 14 | class Split(ModelLayer): |
| 15 | |
| 16 | def __init__(self, model, input_record, num_splits=1, axis=1, |
| 17 | name='split', split=None, **kwargs): |
| 18 | super().__init__(model, name, input_record, **kwargs) |
| 19 | self.axis = axis |
| 20 | # Assume that first dimension is batch, so actual axis in shape is |
| 21 | # axis - 1 |
| 22 | axis -= 1 |
| 23 | assert axis >= 0 |
| 24 | |
| 25 | assert isinstance(input_record, schema.Scalar),\ |
| 26 | "Incorrect input type. Expected Scalar, but received: {0}".\ |
| 27 | format(input_record) |
| 28 | |
| 29 | input_shape = input_record.field_type().shape |
| 30 | assert len(input_shape) >= axis |
| 31 | if split is None: |
| 32 | assert input_shape[axis] % num_splits == 0 |
| 33 | else: |
| 34 | num_splits = len(split) |
| 35 | assert input_shape[axis] == sum(split) |
| 36 | |
| 37 | if split is None: |
| 38 | output_shape = list(input_shape) |
| 39 | output_shape[axis] = int(output_shape[axis] / num_splits) |
| 40 | else: |
| 41 | output_shape = [] |
| 42 | for i in range(num_splits): |
| 43 | output_shape_i = list(input_shape) |
| 44 | output_shape_i[axis] = split[i] |
| 45 | output_shape.append(output_shape_i) |
| 46 | |
| 47 | data_type = input_record.field_type().base |
| 48 | |
| 49 | |
| 50 | if split is None: |
| 51 | output_scalars = [ |
| 52 | schema.Scalar( |
| 53 | (data_type, output_shape), |
| 54 | self.get_next_blob_reference('output_{}'.format(i)), |
| 55 | ) |
| 56 | for i in range(num_splits) |
| 57 | ] |
| 58 | else: |
| 59 | output_scalars = [ |
| 60 | schema.Scalar( |
| 61 | (data_type, output_shape[i]), |
| 62 | self.get_next_blob_reference('output_{}'.format(i)), |
| 63 | ) |
| 64 | for i in range(num_splits) |
| 65 | ] |
| 66 | self.output_schema = schema.Tuple(*output_scalars) |
| 67 | self.split = split |
| 68 | |
| 69 | def add_ops(self, net): |
| 70 | net.Split( |
nothing calls this directly
no test coverage detected