#pragma once
|
|
#include "caffe2/core/operator.h"
|
|
namespace caffe2 {
|
|
struct CAFFE2_API QShapeInfo {
|
QShapeInfo(float o = 0, float s = 1, uint32_t a = 1) {
|
offset.clear();
|
scale.clear();
|
offset.push_back(o);
|
scale.push_back(s);
|
axis = a;
|
}
|
|
uint32_t axis;
|
vector<float> offset;
|
vector<float> scale;
|
};
|
|
struct CAFFE2_API ShapeInfo {
|
enum DimType : int8_t { UNKNOWN = 0, CONSTANT = 1, BATCH = 2, SEQ = 3 };
|
ShapeInfo(bool q = false) : is_quantized(q) {}
|
ShapeInfo(DimType t, TensorShape&& s, bool q = false)
|
: dim_type(t), shape(std::move(s)), is_quantized(q) {}
|
ShapeInfo(DimType t, const TensorShape& s, bool q = false)
|
: dim_type(t), shape(s), is_quantized(q) {}
|
|
ShapeInfo(bool q, const QShapeInfo& info) : is_quantized(q), q_info(info) {}
|
ShapeInfo(DimType t, TensorShape&& s, bool q, const QShapeInfo& info)
|
: dim_type(t), shape(std::move(s)), is_quantized(q), q_info(info) {}
|
ShapeInfo(DimType t, const TensorShape& s, bool q, const QShapeInfo& info)
|
: dim_type(t), shape(s), is_quantized(q), q_info(info) {}
|
|
// type of the shape according its first dim
|
DimType dim_type{DimType::UNKNOWN};
|
TensorShape shape;
|
|
// quantization related information
|
bool is_quantized;
|
QShapeInfo q_info;
|
};
|
|
using ShapeInfoMap = std::unordered_map<std::string, ShapeInfo>;
|
|
// Generates ShapeInfo from Blob.
|
ShapeInfo getShapeInfoFromBlob(const Blob* blob);
|
|
bool operator==(const ShapeInfo& lhs, const ShapeInfo& rhs);
|
|
} // namespace caffe2
|
