/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators/fully_connected_op.h (9351B)
#ifndef CAFFE2_OPERATORS_FULLY_CONNECTED_OP_H_
#define CAFFE2_OPERATORS_FULLY_CONNECTED_OP_H_
#include
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/conversions.h"
#include "caffe2/utils/math.h"
namespace caffe2 {
// This is Caffe's InnerProductOp, with a name that fits its purpose better.
template <
class Context,
class Engine = DefaultEngine,
bool TransposeWeight = true>
class FullyConnectedOp final : public Operator {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
template
explicit FullyConnectedOp(Args&&... args)
: Operator(std::forward(args)...),
axis_(this->template GetSingleArgument("axis", 1)),
axis_w_(this->template GetSingleArgument("axis_w", 1)),
float16_compute_(
this->template GetSingleArgument("float16_compute", false)) {}
~FullyConnectedOp() {}
template <
typename T_X,
typename T_W,
typename T_B,
typename T_Y,
typename MATH>
bool DoRunWithType() {
const auto& X = Input(0);
const auto& W = Input(1);
const auto& b = Input(2);
CAFFE_ENFORCE(b.dim() == 1, b.dim());
// batch size
const auto canonical_axis = X.canonical_axis_index(axis_);
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(axis_w_);
const int N = TransposeWeight ? W.size_to_dim(canonical_axis_w)
: W.size_from_dim(canonical_axis_w);
auto dimErrorString = [&]() {
return c10::str(
"Dimension mismatch: ",
"X: ",
X.sizes(),
", W: ",
W.sizes(),
", b: ",
b.sizes(),
", axis: ",
axis_,
", M: ",
M,
", N: ",
N,
", K: ",
K);
};
// Error checking
CAFFE_ENFORCE(M == X.numel() / K, dimErrorString());
CAFFE_ENFORCE(K == W.numel() / N, dimErrorString());
CAFFE_ENFORCE(N == b.dim32(0), dimErrorString());
CAFFE_ENFORCE(N == b.numel(), dimErrorString());
Y_shape_cache_ = X.sizes().vec();
// This is an invariant of canonical_axis, so we can DCHECK.
DCHECK_LE(canonical_axis + 1, Y_shape_cache_.size());
Y_shape_cache_.resize(canonical_axis + 1);
Y_shape_cache_[canonical_axis] = N;
auto* Y = Output(0, Y_shape_cache_, at::dtype());
CAFFE_ENFORCE(M * N == Y->numel(), dimErrorString());
if (X.numel() == 0) {
// skip the rest of the computation if X is empty
Y->template mutable_data();
return true;
}
// default to FLOAT as math.h does.
TensorProto::DataType math_type = TensorProto_DataType_FLOAT;
if (fp16_type