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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
#ifndef TENSORFLOW_KERNELS_CONV_2D_H_
#define TENSORFLOW_KERNELS_CONV_2D_H_
#include "tensorflow/core/framework/tensor_types.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/NeuralNetworks"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
namespace tensorflow {
namespace functor {
// TODO(yangke): revisit these operations and in particular, see if we can
// combine all of them into just one operation without causing nvcc to
// timeout.
template <typename Device, typename T, int Dims>
struct ShuffleAndReverse {
void operator()(const Device& d, typename TTypes<T, Dims>::ConstTensor input,
const Eigen::DSizes<Eigen::DenseIndex, Dims>& order,
const Eigen::array<bool, Dims>& reverse_dims,
typename TTypes<T, Dims>::Tensor output) {
output.device(d) = input.shuffle(order).reverse(reverse_dims);
}
};
template <typename Device, typename T, int Dims>
struct InflatePadAndShuffle {
void operator()(
const Device& d, typename TTypes<T, Dims>::ConstTensor input,
const Eigen::DSizes<Eigen::DenseIndex, Dims>& strides,
const Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, Dims>& pad_dims,
const Eigen::DSizes<Eigen::DenseIndex, Dims>& order,
typename TTypes<T, Dims>::Tensor output) {
output.device(d) = input.inflate(strides).pad(pad_dims).shuffle(order);
}
};
template <typename Device, typename Input, typename Filter, typename Output>
void SpatialConvolutionFunc(const Device& d, Output output, Input input,
Filter filter, int stride,
const Eigen::PaddingType& padding) {
output.device(d) = Eigen::SpatialConvolution(input, filter, stride, padding);
}
template <typename Device, typename T>
struct SpatialConvolution {
void operator()(const Device& d, typename TTypes<T, 4>::Tensor output,
typename TTypes<T, 4>::ConstTensor input,
typename TTypes<T, 4>::ConstTensor filter, int stride,
const Eigen::PaddingType& padding) {
SpatialConvolutionFunc(d, output, input, filter, stride, padding);
}
};
template <typename Device, typename T>
struct SpatialConvolutionBackwardInput {
void operator()(const Device& d, typename TTypes<T, 4>::Tensor input_backward,
typename TTypes<T, 4>::ConstTensor kernel,
typename TTypes<T, 4>::ConstTensor output_backward,
int input_rows, int input_cols, int stride) {
input_backward.device(d) = Eigen::SpatialConvolutionBackwardInput(
kernel, output_backward, input_rows, input_cols, stride);
}
};
template <typename Device, typename T>
struct SpatialConvolutionBackwardKernel {
void operator()(const Device& d,
typename TTypes<T, 4>::Tensor kernel_backward,
typename TTypes<T, 4>::ConstTensor input,
typename TTypes<T, 4>::ConstTensor output_backward,
int kernel_rows, int kernel_cols, int stride) {
kernel_backward.device(d) = Eigen::SpatialConvolutionBackwardKernel(
input, output_backward, kernel_rows, kernel_cols, stride);
}
};
// TODO(vrv): Figure out how to use the MatMulFunctor in matmul_op.h.
// My initial attempt to do this compiled but failed in the pytest
// due to a swigdeps error.
template <typename Device, typename T>
struct MatMulConvFunctor {
// Computes on device "d": out = in0 * in1, where * is matrix
// multiplication.
void operator()(
const Device& d, typename TTypes<T, 2>::Tensor out,
typename TTypes<T, 2>::ConstTensor in0,
typename TTypes<T, 2>::ConstTensor in1,
const Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1>& dim_pair) {
out.device(d) = in0.contract(in1, dim_pair);
}
};
template <typename Device, typename T>
struct TransformFilter {
void operator()(const Device& d, typename TTypes<T, 4>::ConstTensor in,
typename TTypes<T, 4>::Tensor out) {
out.device(d) = in.shuffle(Eigen::DSizes<Eigen::DenseIndex, 4>(3, 2, 0, 1));
}
};
template <typename Device, typename T>
struct TransformDepth {
void operator()(const Device& d, typename TTypes<T, 4>::ConstTensor in,
const Eigen::DSizes<Eigen::DenseIndex, 4>& shuffle,
typename TTypes<T, 4>::Tensor out) {
out.device(d) = in.shuffle(shuffle);
}
};
template <typename Device, typename T>
struct PadInput {
void operator()(const Device& d, typename TTypes<T, 4>::ConstTensor in,
int padding_rows_left, int padding_rows_right,
int padding_cols_left, int padding_cols_right,
typename TTypes<T, 4>::Tensor out) {
Eigen::array<std::pair<ptrdiff_t, ptrdiff_t>, 4> padding;
padding[0] = std::make_pair(0, 0);
padding[1] = std::make_pair(padding_rows_left, padding_rows_right);
padding[2] = std::make_pair(padding_cols_left, padding_cols_right);
padding[3] = std::make_pair(0, 0);
out.device(d) = in.pad(padding);
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_KERNELS_CONV_2D_H_
|
