1 #include "caffe2/core/context.h" 2 #include "caffe2/core/operator.h" 3 #include "caffe2/operators/conv_pool_op_base.h" 7 #if !EIGEN_VERSION_AT_LEAST(3, 3, 0) 8 #error "Caffe2 requires Eigen to be at least 3.3.0."; 11 #include "unsupported/Eigen/CXX11/Tensor" 21 OPERATOR_NEEDS_FEATURE(group_ == 1,
"Group convolution not supported yet.");
25 bool RunOnDeviceWithOrderNCHW()
override;
26 bool RunOnDeviceWithOrderNHWC()
override;
29 INPUT_TAGS(INPUT, FILTER, BIAS);
36 auto& X = Input(INPUT);
37 auto& filter = Input(FILTER);
39 const int N = X.dim32(0), C = X.dim32(1), H = X.dim32(2), W = X.dim32(3);
40 CAFFE_ENFORCE(4 == filter.ndim());
41 const int M = filter.dim32(0);
42 CAFFE_ENFORCE(filter.dim32(1) == C);
43 CAFFE_ENFORCE(filter.dim32(2) == kernel_h());
44 CAFFE_ENFORCE(filter.dim32(3) == kernel_w());
46 Eigen::array<TIndex, 4> kernel_shuffles
47 { {TIndex(2), TIndex(3), TIndex(1), TIndex(0)} };
48 Eigen::array<TIndex, 4> input_shuffles
49 { {TIndex(0), TIndex(2), TIndex(3), TIndex(1)} };
51 Eigen::Tensor<T, 4, Eigen::RowMajor> filter_tensor =
52 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>>(
53 const_cast<T*
>(filter.template data<T>()),
58 .shuffle(kernel_shuffles);
59 Eigen::Tensor<T, 4, Eigen::RowMajor> X_tensor =
60 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>>(
61 const_cast<T*
>(X.template data<T>()), N, C, H, W)
62 .shuffle(input_shuffles);
67 typedef typename Eigen::internal::traits<
68 Eigen::Tensor<T, 4, Eigen::RowMajor>>
::Index TensorIndex;
69 Eigen::array<Eigen::IndexPair<TensorIndex>, 1> contract_dims;
70 contract_dims[0] = Eigen::IndexPair<TensorIndex>(1, 0);
72 Eigen::DSizes<TensorIndex, 2> pre_contract_dims;
73 pre_contract_dims[1] = kernel_h() * kernel_w() * C;
74 pre_contract_dims[0] = Y->size() / M;
76 Eigen::DSizes<TensorIndex, 2> kernel_dims;
77 kernel_dims[0] = kernel_h() * kernel_w() * C;
80 Eigen::array<TensorIndex, 4> bcast_dims;
82 bcast_dims[1] = Y->dim32(1);
83 bcast_dims[2] = Y->dim32(2);
86 Eigen::Tensor<T, 4, Eigen::RowMajor> Y_tensor(
87 Y->dim32(0), Y->dim32(2), Y->dim32(3), Y->dim32(1));
89 .extract_image_patches(
103 .reshape(pre_contract_dims)
104 .contract(filter_tensor.reshape(kernel_dims), contract_dims)
105 .reshape(Y_tensor.dimensions());
106 if (InputSize() == 3) {
107 auto& bias = Input(BIAS);
108 CAFFE_ENFORCE(1 == bias.ndim());
109 CAFFE_ENFORCE(bias.dim32(0) == M);
112 EigenArrayMap<T> Y_arr(
113 Y_tensor.data(),
static_cast<TIndex
>(M), Y->size() / M);
114 ConstEigenVectorArrayMap<T> bias_arr(bias.template data<T>(), M);
115 Y_arr = Y_arr.colwise() + bias_arr;
119 Eigen::array<TIndex, 4> output_shuffles
120 { {TIndex(0), TIndex(3), TIndex(1), TIndex(2) } };
122 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>>(
123 Y->template mutable_data<T>(), N, M, Y->dim32(2), Y->dim32(3)) =
124 Y_tensor.shuffle(output_shuffles);
128 template <
typename T>
130 auto& X = Input(INPUT);
131 auto& filter = Input(FILTER);
133 const int N = X.dim32(0), H = X.dim32(1), W = X.dim32(2), C = X.dim32(3);
134 CAFFE_ENFORCE(4 == filter.ndim());
135 const int M = filter.dim32(0);
136 CAFFE_ENFORCE(filter.dim32(1) == kernel_h());
137 CAFFE_ENFORCE(filter.dim32(2) == kernel_w());
138 CAFFE_ENFORCE(filter.dim32(3) == C);
142 Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic> temp_filter(
143 M, kernel_h() * kernel_w() * C);
144 temp_filter = ConstEigenArrayMap<T>(
145 filter.template data<T>(), kernel_h() * kernel_w() * C, M)
151 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>> X_tensor(
152 const_cast<T*>(X.template data<T>()), N, H, W, C);
153 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>> Y_tensor(
154 Y->template mutable_data<T>(), N, Y->dim32(1), Y->dim32(2), M);
155 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>> filter_tensor(
156 const_cast<T*>(temp_filter.data()), kernel_h(), kernel_w(), C, M);
161 typedef typename Eigen::internal::traits<
162 Eigen::Tensor<T, 4, Eigen::RowMajor>>
::Index TensorIndex;
163 Eigen::array<Eigen::IndexPair<TensorIndex>, 1> contract_dims;
164 contract_dims[0] = Eigen::IndexPair<TensorIndex>(1, 0);
166 Eigen::DSizes<TensorIndex, 2> pre_contract_dims;
167 pre_contract_dims[1] = kernel_h() * kernel_w() * C;
168 pre_contract_dims[0] = Y->size() / M;
170 Eigen::DSizes<TensorIndex, 2> kernel_dims;
171 kernel_dims[0] = kernel_h() * kernel_w() * C;
174 Eigen::array<TensorIndex, 4> bcast_dims;
176 bcast_dims[1] = Y->dim32(1);
177 bcast_dims[2] = Y->dim32(2);
181 .extract_image_patches(
195 .reshape(pre_contract_dims)
196 .contract(filter_tensor.reshape(kernel_dims), contract_dims)
197 .reshape(Y_tensor.dimensions());
199 if (InputSize() == 3) {
200 auto& bias = Input(BIAS);
201 CAFFE_ENFORCE(1 == bias.ndim());
202 CAFFE_ENFORCE(bias.dim32(0) == M);
203 Eigen::TensorMap<Eigen::Tensor<T, 4, Eigen::RowMajor>> bias_tensor(
204 const_cast<T*>(bias.template data<T>()), 1, 1, 1, M);
207 EigenArrayMap<T> Y_arr(
208 Y->template mutable_data<T>(), static_cast<TIndex>(M), Y->size() / M);
209 ConstEigenVectorArrayMap<T> bias_arr(bias.template data<T>(), M);
210 Y_arr = Y_arr.colwise() + bias_arr;
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...