1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
|
#include "caffe2/operators/locally_connected_op_util.h"
#include <algorithm>
namespace caffe2 {
namespace lc_op_util {
void SetColumnBufferShape(
const int N,
const int kernel_size,
const int output_image_size,
const std::vector<int>& output_image_dims,
const StorageOrder order,
std::vector<int>* column_slice_dims,
std::vector<int>* column_dims,
std::vector<int>* column_transposed_dims,
std::vector<int>* column_axes) {
column_slice_dims->resize(output_image_dims.size() + 1);
if (order == StorageOrder::NCHW) {
column_slice_dims->front() = kernel_size;
std::copy(
output_image_dims.cbegin(),
output_image_dims.cend(),
column_slice_dims->begin() + 1);
} else {
std::copy(
output_image_dims.cbegin(),
output_image_dims.cend(),
column_slice_dims->begin());
column_slice_dims->back() = kernel_size;
}
*column_dims = order == StorageOrder::NCHW
? std::vector<int>{N, kernel_size, output_image_size}
: std::vector<int>{N, output_image_size, kernel_size};
*column_transposed_dims = order == StorageOrder::NCHW
? std::vector<int>{output_image_size, kernel_size, N}
: std::vector<int>{output_image_size, N, kernel_size};
*column_axes = order == StorageOrder::NCHW ? std::vector<int>{2, 1, 0}
: std::vector<int>{1, 0, 2};
}
void SetYBufferShape(
const int N,
const int M,
const int output_image_size,
const StorageOrder order,
std::vector<int>* Y_dims,
std::vector<int>* Y_transposed_dims,
std::vector<int>* Y_axes) {
*Y_dims = order == StorageOrder::NCHW
? std::vector<int>{N, M, output_image_size}
: std::vector<int>{N, output_image_size, M};
*Y_transposed_dims = order == StorageOrder::NCHW
? std::vector<int>{output_image_size, M, N}
: std::vector<int>{output_image_size, N, M};
*Y_axes = order == StorageOrder::NCHW ? std::vector<int>{2, 1, 0}
: std::vector<int>{1, 0, 2};
}
} // namespace lc_op_util
} // namespace caffe2
|
