mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
c10662962c
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17407 As title says Reviewed By: csummersea Differential Revision: D14177921 fbshipit-source-id: e48e1086d37de2c290922d1f498e2d2dad49708a
61 lines
2.2 KiB
C++
61 lines
2.2 KiB
C++
#include "caffe2/operators/fc_inference.h"
|
|
|
|
namespace caffe2 {
|
|
std::vector<TensorShape> FCShapeInference(
|
|
const OperatorDef& def,
|
|
const vector<TensorShape>& in,
|
|
bool pretransposed_weight) {
|
|
vector<TensorShape> out(1);
|
|
|
|
if (in[0].unknown_shape() || in[1].unknown_shape()) {
|
|
out[0].set_unknown_shape(true);
|
|
return out;
|
|
}
|
|
|
|
ArgumentHelper helper(def);
|
|
|
|
auto axis = helper.GetSingleArgument<int32_t>("axis", 1);
|
|
const auto canonical_axis = canonical_axis_index_(axis, in[0].dims().size());
|
|
auto axis_w = helper.GetSingleArgument<int32_t>("axis_w", 1);
|
|
const int canonical_axis_w =
|
|
canonical_axis_index_(axis_w, in[1].dims().size());
|
|
const int64_t N = pretransposed_weight
|
|
? size_from_dim_(canonical_axis_w, GetDimsVector(in[1]))
|
|
: size_to_dim_(canonical_axis_w, GetDimsVector(in[1]));
|
|
|
|
vector<int64_t> y_shape(in[0].dims().begin(), in[0].dims().end());
|
|
CAFFE_ENFORCE_LE(canonical_axis + 1, y_shape.size());
|
|
y_shape.resize(canonical_axis + 1);
|
|
y_shape[canonical_axis] = N;
|
|
|
|
out[0] = CreateTensorShape(y_shape, in[0].data_type());
|
|
return out;
|
|
}
|
|
|
|
OpSchema::Cost CostInferenceForFC(
|
|
const OperatorDef& def,
|
|
const vector<TensorShape>& in,
|
|
bool pretransposed_weight) {
|
|
CAFFE_ENFORCE_EQ(in.size(), 3, "FC requires three inputs");
|
|
struct OpSchema::Cost c;
|
|
ArgumentHelper helper(def);
|
|
|
|
auto axis = helper.GetSingleArgument<int32_t>("axis", 1);
|
|
const auto canonical_axis = canonical_axis_index_(axis, in[0].dims().size());
|
|
const uint64_t M = size_to_dim_(canonical_axis, GetDimsVector(in[0]));
|
|
const uint64_t K = size_from_dim_(canonical_axis, GetDimsVector(in[0]));
|
|
auto axis_w = helper.GetSingleArgument<int32_t>("axis_w", 1);
|
|
const int canonical_axis_w =
|
|
canonical_axis_index_(axis_w, in[1].dims().size());
|
|
const uint64_t N = pretransposed_weight
|
|
? size_from_dim_(canonical_axis_w, GetDimsVector(in[1]))
|
|
: size_to_dim_(canonical_axis_w, GetDimsVector(in[1]));
|
|
|
|
const auto& X = in[0];
|
|
c.flops = M * N * (2 * K + 1);
|
|
c.bytes_read = (K * (M + N) + N) * sizeof(X.data_type());
|
|
c.bytes_written = M * N * sizeof(X.data_type());
|
|
c.params_bytes = (K * N + N) * sizeof(X.data_type());
|
|
return c;
|
|
}
|
|
} // namespace caffe2
|