summaryrefslogtreecommitdiff
path: root/caffe2
diff options
context:
space:
mode:
Diffstat (limited to 'caffe2')
-rw-r--r--caffe2/operators/utility_ops.cc38
-rw-r--r--caffe2/operators/utility_ops.h12
-rw-r--r--caffe2/python/operator_test/flatten_op_test.py38
-rw-r--r--caffe2/python/operator_test/shape_inference_test.py16
4 files changed, 86 insertions, 18 deletions
diff --git a/caffe2/operators/utility_ops.cc b/caffe2/operators/utility_ops.cc
index 5763e64a35..9a2e82e09f 100644
--- a/caffe2/operators/utility_ops.cc
+++ b/caffe2/operators/utility_ops.cc
@@ -135,34 +135,44 @@ OPERATOR_SCHEMA(LengthsToShape).NumInputs(1).NumOutputs(1);
OPERATOR_SCHEMA(Flatten)
.NumInputs(1)
.NumOutputs(1)
- .TensorInferenceFunction([](const OperatorDef&,
+ .TensorInferenceFunction([](const OperatorDef& def,
const vector<TensorShape>& in) {
+ ArgumentHelper helper(def);
+ const int axis = helper.GetSingleArgument<int>("axis", 1);
vector<TensorShape> out(1);
- int total = 1;
+ TIndex outer = 1;
+ TIndex inner = 1;
std::size_t index = 0;
for (auto d : in[0].dims()) {
- // skip the first element
- if (index++ == 0) {
- continue;
+ if (index < axis) {
+ outer *= d;
+ } else {
+ inner *= d;
}
- total *= d;
+ ++index;
}
out[0].set_data_type(in[0].data_type());
- out[0].add_dims(in[0].dims(0));
- out[0].add_dims(total);
+ out[0].add_dims(outer);
+ out[0].add_dims(inner);
return out;
})
.SetDoc(R"DOC(
-Flattens the input tensor into a 2D matrix, keeping the first dimension
-unchanged.
+Flattens the input tensor into a 2D matrix. If input tensor has shape
+(d_0, d_1, ... d_n) then the output will have shape
+(d_0 X d_1 ... d_(axis-1), d_axis X d_(axis+1) ... X dn)
)DOC")
- .Input(0, "input", "A tensor of rank >= 2.")
+ .Input(0, "input", "A tensor of rank >= axis.")
.Output(
0,
"output",
- "A tensor of rank 2 with the contents of the input tensor, "
- "with first dimension equal first dimension of input, and remaining "
- "input dimensions flattened into the inner dimension of the output.");
+ "A 2D tensor with the contents of the input tensor, "
+ "with input dimensions up to axis flattened to the outer dimension "
+ "of the output and remaining input dimensions flattened into the inner "
+ "dimension of the output.")
+ .Arg(
+ "axis",
+ "(Default to 1) Indicate up to which input dimensions "
+ "(exclusive) should be flattened to the outer dimension of the output");
OPERATOR_SCHEMA(FlattenToVec)
.NumInputs(1)
diff --git a/caffe2/operators/utility_ops.h b/caffe2/operators/utility_ops.h
index da080a6cfc..0b130d9655 100644
--- a/caffe2/operators/utility_ops.h
+++ b/caffe2/operators/utility_ops.h
@@ -222,14 +222,17 @@ template <class Context>
class FlattenOp : public Operator<Context> {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
- USE_SIMPLE_CTOR_DTOR(FlattenOp);
+
+ FlattenOp(const OperatorDef& operator_def, Workspace* ws)
+ : Operator<Context>(operator_def, ws),
+ axis_(OperatorBase::GetSingleArgument<int>("axis", 1)) {}
bool RunOnDevice() override {
auto& input = Input(0);
auto* output = Output(0);
CAFFE_ENFORCE_GE(
- input.dims().size(), 2, "The rank of the tensor must be >= 2.");
- output->Resize(input.dim(0), input.size_from_dim(1));
+ input.dims().size(), axis_, "The rank of the tensor must be >= axis.");
+ output->Resize(input.size_to_dim(axis_), input.size_from_dim(axis_));
context_.template CopyItems<Context, Context>(
input.meta(),
input.size(),
@@ -237,6 +240,9 @@ class FlattenOp : public Operator<Context> {
output->raw_mutable_data(input.meta()));
return true;
}
+
+ private:
+ int axis_;
};
template <class Context>
diff --git a/caffe2/python/operator_test/flatten_op_test.py b/caffe2/python/operator_test/flatten_op_test.py
new file mode 100644
index 0000000000..19d204e0bd
--- /dev/null
+++ b/caffe2/python/operator_test/flatten_op_test.py
@@ -0,0 +1,38 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from hypothesis import given
+import numpy as np
+
+from caffe2.python import core
+import caffe2.python.hypothesis_test_util as hu
+
+
+class TestFlatten(hu.HypothesisTestCase):
+ @given(X=hu.tensor(min_dim=2, max_dim=4),
+ **hu.gcs)
+ def test_flatten(self, X, gc, dc):
+ for axis in range(X.ndim + 1):
+ op = core.CreateOperator(
+ "Flatten",
+ ["X"],
+ ["Y"],
+ axis=axis)
+
+ def flatten_ref(X):
+ shape = X.shape
+ outer = np.prod(shape[:axis]).astype(int)
+ inner = np.prod(shape[axis:]).astype(int)
+ return np.copy(X).reshape(outer, inner),
+
+ self.assertReferenceChecks(gc, op, [X], flatten_ref)
+
+ # Check over multiple devices
+ self.assertDeviceChecks(dc, op, [X], [0])
+
+
+if __name__ == "__main__":
+ import unittest
+ unittest.main()
diff --git a/caffe2/python/operator_test/shape_inference_test.py b/caffe2/python/operator_test/shape_inference_test.py
index db7520a179..e903e8ef2e 100644
--- a/caffe2/python/operator_test/shape_inference_test.py
+++ b/caffe2/python/operator_test/shape_inference_test.py
@@ -347,7 +347,7 @@ class TestShapeInference(test_util.TestCase):
self.InferTensorRunAndCompare(model)
-
+ # test Flatten with default axis (=1)
model = model_helper.ModelHelper(name="test_model")
model.Flatten("X", "Flat")
model.Flatten("empty", "EmptyFlat")
@@ -356,6 +356,20 @@ class TestShapeInference(test_util.TestCase):
self.InferTensorRunAndCompare(model)
+ # test Flatten with axis
+ model = model_helper.ModelHelper(name="test_model")
+ x = np.random.randn(17, 5, 13)
+ for axis in range(x.ndim + 1):
+ model.Flatten("x", "Flat", axis=axis)
+ workspace.FeedBlob("x", x)
+ self.InferTensorRunAndCompare(model)
+
+ empty = np.random.randn(0, 5, 13)
+ for axis in range(empty.ndim + 1):
+ model.Flatten("empty", "Flat", axis=axis)
+ workspace.FeedBlob("empty", empty)
+ self.InferTensorRunAndCompare(model)
+
def testShapeInferenceReshape(self):
model = model_helper.ModelHelper(name="test_model")
model.Reshape("X", ["Reshaped", "Old_Shape"], shape=[8, 0, -1, 2])
generated by cgit v1.2.3 (git 2.25.1) at 2024-11-27 07:36:23 +0000