#include <algorithm>
|
#include <vector>
|
#include "caffe2/core/tensor.h"
|
#include "caffe2/utils/eigen_utils.h"
|
#include "caffe2/utils/math.h"
|
|
namespace caffe2 {
|
namespace {
|
|
using t_tuple = std::tuple<Tensor, Tensor>;
|
|
template <typename T>
|
T copy_ctor(const T& x) {
|
return x;
|
}
|
|
template <>
|
Tensor copy_ctor(const Tensor& X) {
|
return X.UnsafeSharedInstance();
|
}
|
|
template <>
|
t_tuple copy_ctor(const t_tuple& X) {
|
return std::make_tuple(copy_ctor(std::get<0>(X)), copy_ctor(std::get<1>(X)));
|
}
|
|
template <>
|
std::pair<t_tuple, t_tuple> copy_ctor(const std::pair<t_tuple, t_tuple>& X) {
|
return std::make_pair(copy_ctor(X.first), copy_ctor(X.second));
|
}
|
|
template <>
|
std::vector<Tensor> copy_ctor(const std::vector<Tensor>& X) {
|
std::vector<Tensor> Y(X.size());
|
std::transform(X.begin(), X.end(), Y.begin(), [](const Tensor& x) {
|
return copy_ctor(x);
|
});
|
return Y;
|
}
|
|
template <>
|
std::vector<t_tuple> copy_ctor(const std::vector<t_tuple>& X) {
|
std::vector<t_tuple> Y(X.size());
|
std::transform(X.begin(), X.end(), Y.begin(), [](const t_tuple& x) {
|
return copy_ctor(x);
|
});
|
return Y;
|
}
|
|
template <>
|
std::vector<std::pair<t_tuple, t_tuple>> copy_ctor(
|
const std::vector<std::pair<t_tuple, t_tuple>>& X) {
|
std::vector<std::pair<t_tuple, t_tuple>> Y(X.size());
|
std::transform(
|
X.begin(), X.end(), Y.begin(), [](const std::pair<t_tuple, t_tuple>& x) {
|
return copy_ctor(x);
|
});
|
return Y;
|
}
|
|
// Gathers every two elements of a vector in a vector of pairs
|
template <typename T>
|
static std::vector<std::pair<T, T>> pair_vec(const std::vector<T>& vals) {
|
CAFFE_ENFORCE_EQ(
|
vals.size() % 2,
|
0,
|
"Odd number of params or hiddens given to a bidirectional RNN");
|
std::vector<std::pair<T, T>> result;
|
result.reserve(vals.size() / 2);
|
for (int64_t i = 0; i < vals.size(); i += 2) {
|
result.emplace_back(copy_ctor(vals[i]), copy_ctor(vals[i + 1]));
|
}
|
return result;
|
}
|
|
// Flattens a vector of pairs
|
template <typename T>
|
static std::vector<T> unpair_vec(std::vector<std::pair<T, T>>&& vals) {
|
std::vector<T> result;
|
result.reserve(vals.size() * 2);
|
for (int64_t i = 0; i < vals.size(); i++) {
|
result.push_back(std::move(vals[i].first));
|
result.push_back(std::move(vals[i].second));
|
}
|
return result;
|
}
|
|
Tensor matmul(const Tensor& X, const Tensor& W, CPUContext* context) {
|
const auto canonical_axis = X.canonical_axis_index(1);
|
const auto M = X.size_to_dim(canonical_axis);
|
const auto K = X.size_from_dim(canonical_axis);
|
const auto canonical_axis_w = W.canonical_axis_index(1);
|
const int N = W.size_to_dim(canonical_axis_w);
|
auto output_size = X.sizes().vec();
|
output_size.resize(canonical_axis + 1);
|
output_size[canonical_axis] = N;
|
Tensor C(output_size, CPU);
|
math::Gemm<float, CPUContext>(
|
CblasNoTrans,
|
CblasTrans,
|
M,
|
N,
|
K,
|
1,
|
X.template data<float>(),
|
W.template data<float>(),
|
0,
|
C.template mutable_data<float>(),
|
context);
|
return C;
|
}
|
|
Tensor
|
linear(const Tensor& X, const Tensor& W, const Tensor& B, CPUContext* context) {
|
auto output = matmul(X, W, context);
|
if (B) {
|
const auto canonical_axis = X.canonical_axis_index(1);
|
const auto M = X.size_to_dim(canonical_axis);
|
const auto canonical_axis_w = W.canonical_axis_index(1);
|
const int N = W.size_to_dim(canonical_axis_w);
|
auto bias_multiplier_ = caffe2::empty({M}, CPU);
|
math::Set<float, CPUContext>(
|
M, 1, bias_multiplier_.template mutable_data<float>(), context);
|
math::Gemm<float, CPUContext>(
|
CblasNoTrans,
|
CblasNoTrans,
|
M,
|
N,
|
1,
|
1,
|
bias_multiplier_.template data<float>(),
|
B.template data<float>(),
|
1,
|
output.template mutable_data<float>(),
|
context);
|
}
|
return output;
|
}
|
|
std::vector<Tensor>
|
chunk(const Tensor& input, int chunks, int axis, CPUContext* context) {
|
int canonical_axis = input.canonical_axis_index(axis);
|
CAFFE_ENFORCE_LT(
|
canonical_axis, input.dim(), "Axis not in input ndim range.");
|
const int input_channels = input.dim32(canonical_axis);
|
CAFFE_ENFORCE_EQ(
|
input_channels % chunks,
|
0,
|
"input channels should be divisible by the number of chunks.");
|
auto split_size = input_channels / chunks;
|
vector<int64_t> output_dims(input.sizes().vec());
|
int before = 1, after = 1;
|
for (int i = 0; i < canonical_axis; ++i) {
|
before *= input.dim32(i);
|
}
|
for (int i = canonical_axis + 1; i < input.dim(); ++i) {
|
after *= input.dim32(i);
|
}
|
size_t input_offset = 0;
|
std::vector<Tensor> outputs;
|
for (int i = 0; i < chunks; ++i) {
|
auto axis_dim = split_size;
|
output_dims[canonical_axis] = split_size;
|
Tensor output(output_dims, CPU);
|
math::CopyMatrix<CPUContext>(
|
input.itemsize(),
|
before,
|
axis_dim * after,
|
static_cast<const char*>(input.raw_data()) + input_offset,
|
input.dim32(canonical_axis) * after,
|
output.raw_mutable_data(input.dtype()),
|
axis_dim * after,
|
context,
|
input.dtype().copy());
|
input_offset += axis_dim * after * input.itemsize();
|
outputs.push_back(std::move(output));
|
}
|
return outputs;
|
}
|
|
std::vector<Tensor> unbind(const Tensor& input, int axis, CPUContext* context) {
|
// 1 - Chunk the input tensor along the given axis into N chunks where
|
// N is the dim(axis)
|
auto chunks = chunk(input, input.sizes()[axis], axis, context);
|
// 2 - Compute new dimensions
|
std::vector<int64_t> newDims = input.sizes().vec();
|
newDims.erase(newDims.begin() + axis);
|
|
// 3 - Reshape chunks to drop the extra dimension
|
for (int i = 0; i < chunks.size(); i++) {
|
CAFFE_ENFORCE_EQ(
|
chunks[i].sizes()[axis], 1, "Got an unexpected chunk size");
|
chunks[i].Reshape(newDims);
|
}
|
return chunks;
|
}
|
|
Tensor
|
cat(const std::vector<Tensor>& tensorList, int axis, CPUContext* context) {
|
// Adopted from C2's concat operator
|
auto input_zero = copy_ctor(tensorList.at(0));
|
vector<int64_t> outputDims(input_zero.sizes().vec());
|
CAFFE_ENFORCE(outputDims.size() > 0);
|
for (int i = 1; i < tensorList.size(); i++) {
|
CAFFE_ENFORCE(input_zero.dtype() == tensorList.at(i).dtype());
|
outputDims[axis] += tensorList.at(i).sizes()[axis];
|
}
|
auto output_channels = outputDims[axis];
|
Tensor output(outputDims, CPU);
|
int before = 1, after = 1;
|
for (int i = 0; i < tensorList.at(0).dim(); ++i) {
|
if (i == axis) {
|
continue;
|
}
|
int dim = input_zero.dim32(i);
|
if (i < axis) {
|
before *= dim;
|
} else {
|
after *= dim;
|
}
|
}
|
size_t output_offset = 0;
|
for (const auto& input : tensorList) {
|
auto axis_dim = input.dim32(axis);
|
math::CopyMatrix<CPUContext>(
|
input.itemsize(),
|
before,
|
axis_dim * after,
|
input.raw_data(),
|
axis_dim * after,
|
static_cast<char*>(output.raw_mutable_data(input_zero.dtype())) +
|
output_offset,
|
output_channels * after,
|
context,
|
input_zero.dtype().copy());
|
output_offset += axis_dim * after * input.itemsize();
|
}
|
|
return output;
|
}
|
|
Tensor
|
stack(const std::vector<Tensor>& tensorList, int axis, CPUContext* context) {
|
// 1 - Compute new dimensions
|
std::vector<int64_t> newDims(tensorList[0].sizes().vec());
|
std::vector<Tensor> expandedTensorList;
|
newDims.insert(newDims.begin() + axis, 1);
|
for (int i = 0; i < tensorList.size(); i++) {
|
expandedTensorList.emplace_back(tensorList[i].Clone());
|
expandedTensorList.at(i).Reshape(newDims);
|
}
|
return cat(expandedTensorList, axis, context);
|
}
|
|
Tensor sigmoid(const Tensor& X) {
|
Tensor Y(X.sizes(), CPU);
|
auto N = X.numel();
|
EigenVectorArrayMap<float>(Y.template mutable_data<float>(), N) = 1.0 /
|
(1.0 +
|
(-ConstEigenVectorArrayMap<float>(X.template data<float>(), N)).exp());
|
return Y;
|
}
|
|
Tensor tanh(const Tensor& X, CPUContext* context) {
|
Tensor Y(X.sizes(), CPU);
|
math::Tanh<float, CPUContext>(
|
X.numel(),
|
X.template data<float>(),
|
Y.template mutable_data<float>(),
|
context);
|
return Y;
|
}
|
|
Tensor add(const Tensor& X, const Tensor& Y, CPUContext* context) {
|
Tensor Z(X.sizes().vec(), CPU);
|
math::Add<float, CPUContext>(
|
X.numel(),
|
X.template data<float>(),
|
Y.template data<float>(),
|
Z.template mutable_data<float>(),
|
context);
|
return Z;
|
}
|
|
Tensor mul(const Tensor& X, const Tensor& Y, CPUContext* context) {
|
Tensor Z(X.sizes().vec(), CPU);
|
math::Mul<float, CPUContext>(
|
X.numel(),
|
X.template data<float>(),
|
Y.template data<float>(),
|
Z.template mutable_data<float>(),
|
context);
|
return Z;
|
}
|
|
Tensor transpose(const Tensor& X, int dim0, int dim1, CPUContext* context) {
|
int ndim = X.dim();
|
CAFFE_ENFORCE(ndim > dim0 && ndim > dim1, "Invalid transpose dimensions");
|
std::vector<int> axes(ndim);
|
std::iota(axes.begin(), axes.end(), 0);
|
std::swap(axes[dim0], axes[dim1]);
|
const std::vector<std::int64_t> X_dims = X.sizes().vec();
|
std::vector<std::int64_t> Y_dims(ndim);
|
for (int i = 0; i < ndim; ++i) {
|
Y_dims[i] = X_dims[axes[i]];
|
}
|
Tensor Y(Y_dims, CPU);
|
math::Transpose<std::int64_t, float, CPUContext>(
|
ndim,
|
X_dims.data(),
|
axes.data(),
|
X.template data<float>(),
|
Y.template mutable_data<float>(),
|
context);
|
return Y;
|
}
|
} // namespace
|
} // namespace caffe2
|
