提交 71de31e0 编写于 作者: L luxuhui

fix: fix bugs on security

N/A
Signed-off-by: NLuxuhui <luxuhui@xiaomi.com>
上级 a455ca53
......@@ -416,7 +416,8 @@ class MACE_API MaceEngine {
MaceStatus Init(const NetDef *net_def,
const std::vector<std::string> &input_nodes,
const std::vector<std::string> &output_nodes,
const unsigned char *model_data);
const unsigned char *model_data,
const int64_t model_data_size);
MaceStatus Init(const NetDef *net_def,
const std::vector<std::string> &input_nodes,
......
......@@ -74,11 +74,8 @@ MACE_GET_OPTIONAL_ARGUMENT_FUNC(std::string, s, false)
#define MACE_GET_REPEATED_ARGUMENT_FUNC(T, fieldname, lossless_conversion) \
template <> \
std::vector<T> ProtoArgHelper::GetRepeatedArgs<T>( \
const std::string &arg_name, const std::vector<T> &default_value) \
const { \
if (arg_map_.count(arg_name) == 0) { \
return default_value; \
} \
const std::string &arg_name) const { \
MACE_CHECK(arg_map_.count(arg_name) > 0, arg_name, "not exist."); \
std::vector<T> values; \
for (const auto &v : arg_map_.at(arg_name).fieldname()) { \
if (lossless_conversion) { \
......@@ -89,8 +86,19 @@ MACE_GET_OPTIONAL_ARGUMENT_FUNC(std::string, s, false)
values.push_back(v); \
} \
return values; \
} \
template <> \
std::vector<T> ProtoArgHelper::GetRepeatedArgs<T>( \
const std::string &arg_name, const std::vector<T> &default_value) \
const { \
if (arg_map_.count(arg_name) == 0) { \
return default_value; \
} else { \
return GetRepeatedArgs<T>(arg_name); \
} \
}
MACE_GET_REPEATED_ARGUMENT_FUNC(float, floats, false)
MACE_GET_REPEATED_ARGUMENT_FUNC(int, ints, true)
MACE_GET_REPEATED_ARGUMENT_FUNC(int64_t, ints, true)
......
......@@ -46,10 +46,14 @@ class ProtoArgHelper {
template <typename T>
T GetOptionalArg(const std::string &arg_name, const T &default_value) const;
template <typename T>
std::vector<T> GetRepeatedArgs(
const std::string &arg_name,
const std::vector<T> &default_value = std::vector<T>()) const;
const std::vector<T> &default_value) const;
template <typename T>
std::vector<T> GetRepeatedArgs(const std::string &arg_name) const;
private:
std::map<std::string, Argument> arg_map_;
......
......@@ -114,12 +114,9 @@ void MemoryOptimizer::Optimize(
const mace::OperatorDef *op_def,
const std::unordered_map<std::string, MemoryType> *mem_types) {
MACE_LATENCY_LOGGER(2, "Optimize memory");
if (op_def->output_size() != op_def->output_shape_size()) {
VLOG(1) << op_def->name()
<< ": the number of output shape "
<< "is not equal to the number of output";
return;
}
MACE_CHECK(op_def->output_size() == op_def->output_shape_size(),
op_def->name(), "The number of output shapes is"
" not equal to the number of outputs");
auto device = static_cast<DeviceType>(op_def->device_type());
DataType op_dtype = static_cast<DataType>(ProtoArgHelper::GetOptionalArg(
......@@ -220,8 +217,8 @@ void MemoryOptimizer::Optimize(
} else {
mem_ref_count_[best_mem_id] = 1;
}
tensor_mem_map_.emplace(op_def->output(i), TensorMemInfo(best_mem_id,
dt, data_format));
tensor_mem_map_.emplace(op_def->output(i),
TensorMemInfo(best_mem_id, dt, data_format));
}
}
......@@ -249,7 +246,7 @@ void MemoryOptimizer::Optimize(
}
}
const std::vector<MemoryBlock>& MemoryOptimizer::mem_blocks() const {
const std::vector<MemoryBlock> &MemoryOptimizer::mem_blocks() const {
return mem_blocks_;
}
......@@ -277,8 +274,8 @@ std::string MemoryOptimizer::DebugInfo() const {
sstream << i << " " << memory_type_to_str(mem_blocks_[i].mem_type())
<< " ";
if (mem_blocks_[i].mem_type() == MemoryType::GPU_IMAGE) {
sstream << DataTypeToString(mem_blocks_[i].data_type()) << " "
"[" << mem_blocks_[i].x() << ", " << mem_blocks_[i].y() << "]";
sstream << DataTypeToString(mem_blocks_[i].data_type()) << " ["
<< mem_blocks_[i].x() << ", " << mem_blocks_[i].y() << "]";
} else {
sstream << "[" << mem_blocks_[i].x() << "]";
}
......
......@@ -94,9 +94,9 @@ class MemoryOptimizer {
const MemoryType mem_type);
void UpdateTensorRef(const std::string &tensor_name);
void UpdateTensorRef(const OperatorDef *op_def);
void Optimize(
virtual void Optimize(
const OperatorDef *op_def,
const std::unordered_map<std::string, MemoryType> *mem_types = nullptr);
const std::unordered_map<std::string, MemoryType> *mem_types);
const std::vector<MemoryBlock> &mem_blocks() const;
......
......@@ -54,6 +54,9 @@ SerialNet::SerialNet(const OpRegistry *op_registry,
#ifdef MACE_ENABLE_OPENCL
// used for memory optimization
std::unordered_map<std::string, MemoryType> output_mem_map;
std::unordered_map<std::string, MemoryType> *mem_types = &output_mem_map;
#else
std::unordered_map<std::string, MemoryType> *mem_types = nullptr;
#endif // MACE_ENABLE_OPENCL
OpConstructContext construct_context(ws_);
......@@ -100,11 +103,7 @@ SerialNet::SerialNet(const OpRegistry *op_registry,
for (auto &op : operators_) {
VLOG(2) << "Operator " << op->debug_def().name() << "<" << op->device_type()
<< ", " << op->debug_def().type() << ">";
#ifdef MACE_ENABLE_OPENCL
mem_optimizer->Optimize(op->operator_def().get(), &output_mem_map);
#else
mem_optimizer->Optimize(op->operator_def().get());
#endif // MACE_ENABLE_OPENCL
mem_optimizer->Optimize(op->operator_def().get(), mem_types);
}
VLOG(1) << mem_optimizer->DebugInfo();
}
......@@ -128,8 +127,7 @@ MaceStatus SerialNet::Init() {
MaceStatus SerialNet::Run(RunMetadata *run_metadata) {
const char *profiling = getenv("MACE_OPENCL_PROFILING");
bool
enable_opencl_profiling =
bool enable_opencl_profiling =
profiling != nullptr && strlen(profiling) == 1 && profiling[0] == '1';
MACE_MEMORY_LOGGING_GUARD();
......
......@@ -48,14 +48,22 @@ class Operation {
return ProtoArgHelper::GetOptionalArg<OperatorDef, T>(
*operator_def_, name, default_value);
}
template<typename T>
std::vector<T> GetRepeatedArgs(
const std::string &name, const std::vector<T> &default_value = {}) const {
const std::string &name, const std::vector<T> &default_value) const {
MACE_CHECK(operator_def_, "operator_def was null!");
return ProtoArgHelper::GetRepeatedArgs<OperatorDef, T>(
*operator_def_, name, default_value);
}
template<typename T>
std::vector<T> GetRepeatedArgs(const std::string &name) const {
MACE_CHECK(operator_def_, "operator_def was null!");
return ProtoArgHelper::GetRepeatedArgs<OperatorDef, T>(
*operator_def_, name);
}
DeviceType device_type() const {
return static_cast<DeviceType>(operator_def_->device_type());
}
......@@ -65,7 +73,10 @@ class Operation {
return inputs_[idx];
}
Tensor *Output(int idx) { return outputs_[idx]; }
Tensor *Output(unsigned int idx) {
MACE_CHECK(idx < outputs_.size());
return outputs_[idx];
}
int InputSize() { return inputs_.size(); }
int OutputSize() { return outputs_.size(); }
......
......@@ -74,7 +74,8 @@ apu_eltwise_mode ApuWrapper::MapToApuEltwiseMode(int mace_mode) {
}
bool ApuWrapper::Init(const NetDef &net_def,
unsigned const char *model_data) {
unsigned const char *model_data,
const index_t model_data_size) {
frontend = new ApuFrontend();
// parse model argument
......@@ -100,11 +101,15 @@ bool ApuWrapper::Init(const NetDef &net_def,
tensor.dim_size = const_tensor.dims_size();
MACE_CHECK(tensor.dim_size <= APU_TENSOR_MAX_DIMS,
"tensor dimension size not supported");
for (auto i = 0 ; i < tensor.dim_size ; i++) {
for (auto i = 0; i < tensor.dim_size; i++) {
tensor.dims[i] = const_tensor.dims(i);
}
const auto tensor_end = const_tensor.offset() +
const_tensor->data_size() * GetEnumTypeSize(const_tensor.data_type());
MACE_CHECK(tensor_end <= model_data_size, "tensor_end (", tensor_end,
") should <= ", model_data_size);
tensor.data_buf =
const_cast<unsigned char*>(model_data + const_tensor.offset());
const_cast<unsigned char *>(model_data + const_tensor.offset());
const_tensors.push_back(tensor);
}
......@@ -121,10 +126,10 @@ bool ApuWrapper::Init(const NetDef &net_def,
tensor.dim_size = input_info.dims_size();
MACE_CHECK(tensor.dim_size <= APU_TENSOR_MAX_DIMS,
"tensor dimension size not supported");
tensor_info info;
ApuTensorInfo info;
info.name = input_info.name();
info.size = 1;
for (auto i = 0 ; i < tensor.dim_size ; i++) {
for (auto i = 0; i < tensor.dim_size; i++) {
tensor.dims[i] = input_info.dims(i);
info.size *= input_info.dims(i);
info.shape.push_back(input_info.dims(i));
......@@ -140,13 +145,13 @@ bool ApuWrapper::Init(const NetDef &net_def,
// output tensors
std::vector<int> output_tensor_ids;
std::vector<void*> output_buffers;
std::vector<void *> output_buffers;
for (auto output_info : net_def.output_info()) {
output_tensor_ids.push_back(output_info.node_id());
tensor_info info;
ApuTensorInfo info;
info.name = output_info.name();
info.size = 1;
for (auto i = 0 ; i < output_info.dims().size() ; i++) {
for (auto i = 0; i < output_info.dims().size(); i++) {
info.size *= output_info.dims(i);
info.shape.push_back(output_info.dims(i));
}
......@@ -170,7 +175,7 @@ bool ApuWrapper::Init(const NetDef &net_def,
strncpy(op.type, op_def.type().c_str(), APU_OP_TYPE_MAX_SIZE);
op.input_size = op_def.node_input_size();
std::vector<int> input_ids;
for (auto i = 0 ; i < op.input_size ; i++) {
for (auto i = 0; i < op.input_size; i++) {
input_ids.push_back(op_def.node_input(i).node_id());
}
cached_op_inputs.push_back(input_ids);
......@@ -188,7 +193,7 @@ bool ApuWrapper::Init(const NetDef &net_def,
op.output.dim_size = op_def.output_shape(0).dims_size();
MACE_CHECK(op.output.dim_size <= APU_TENSOR_MAX_DIMS,
"tensor dimension size not supported");
for (auto i = 0 ; i < op.output.dim_size ; i++) {
for (auto i = 0; i < op.output.dim_size; i++) {
op.output.dims[i] = op_def.output_shape(0).dims(i);
}
op.output.data_buf = nullptr;
......@@ -223,12 +228,12 @@ bool ApuWrapper::Init(const NetDef &net_def,
bool print_model = false;
bool ret = frontend->InitGraph(
const_tensors.size(), const_tensors.data(),
input_tensors.size(), input_tensors.data(),
output_tensor_ids.size(), output_tensor_ids.data(),
output_buffers.data(),
ops.size(), ops.data(),
print_model);
const_tensors.size(), const_tensors.data(),
input_tensors.size(), input_tensors.data(),
output_tensor_ids.size(), output_tensor_ids.data(),
output_buffers.data(),
ops.size(), ops.data(),
print_model);
cached_op_inputs.clear();
MACE_CHECK(ret == true, "apu init graph failed");
......@@ -243,7 +248,7 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
// prepare input
for (int i = 0 ; i < static_cast<int>(input_tensors.size()) ; i++) {
Tensor* tensor = input_tensors.at(input_infos[i].name);
Tensor *tensor = input_tensors.at(input_infos[i].name);
// check size
int size = input_infos[i].size;
......@@ -251,7 +256,7 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
// quantize
quantize_util_.QuantizeWithScaleAndZeropoint(
(const float*)tensor->raw_data(),
tensor->data<float>(),
size,
input_infos[i].scale,
input_infos[i].zero_point,
......@@ -263,8 +268,8 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
MACE_CHECK(ret == true, "neuron run model failed");
// process output
for (int i = 0 ; i < static_cast<int>(output_tensors->size()) ; i++) {
Tensor* tensor = output_tensors->at(output_infos[i].name);
for (int i = 0; i < static_cast<int>(output_tensors->size()); i++) {
Tensor *tensor = output_tensors->at(output_infos[i].name);
// prepare out buffer
tensor->SetDtype(DT_FLOAT);
......@@ -278,18 +283,18 @@ bool ApuWrapper::Run(const std::map<std::string, Tensor *> &input_tensors,
size,
output_infos[i].scale,
output_infos[i].zero_point,
reinterpret_cast<float*>(tensor->raw_mutable_data()));
tensor->mutable_data<float>());
}
return true;
}
bool ApuWrapper::Uninit() {
bool ret = frontend->UninitGraph();
frontend = nullptr;
input_infos.clear();
output_infos.clear();
return ret;
bool ret = frontend->UninitGraph();
frontend = nullptr;
input_infos.clear();
output_infos.clear();
return ret;
}
} // namespace mace
......@@ -21,41 +21,42 @@
#include <memory>
#include "mace/proto/mace.pb.h"
#include "mace/core/tensor.h"
#include "mace/core/device.h"
#include "mace/core/quantize.h"
#include "mace/core/tensor.h"
#include "third_party/apu/ApuFrontend.h"
namespace mace {
class ApuWrapper {
struct tensor_info {
struct ApuTensorInfo {
std::string name;
std::shared_ptr<uint8_t> buf;
std::vector<index_t> shape;
int size;
float scale;
int zero_point;
};
};
public:
explicit ApuWrapper(Device *device);
bool Init(const NetDef& net_def, unsigned const char *model_data);
bool Run(const std::map<std::string, Tensor *> &input_tensors,
std::map<std::string, Tensor *> *output_tensors);
bool Uninit();
explicit ApuWrapper(Device *device);
bool Init(const NetDef &net_def, unsigned const char *model_data,
const index_t model_data_size);
bool Run(const std::map<std::string, Tensor *> &input_tensors,
std::map<std::string, Tensor *> *output_tensors);
bool Uninit();
private:
apu_data_type MapToApuDataType(DataType mace_type);
apu_pooling_mode MapToApuPoolingMode(int mace_mode);
apu_eltwise_mode MapToApuEltwiseMode(int mace_mode);
apu_data_type MapToApuDataType(DataType mace_type);
apu_pooling_mode MapToApuPoolingMode(int mace_mode);
apu_eltwise_mode MapToApuEltwiseMode(int mace_mode);
private:
ApuFrontend* frontend;
std::vector<tensor_info> input_infos;
std::vector<tensor_info> output_infos;
QuantizeUtil<float, uint8_t> quantize_util_;
ApuFrontend *frontend;
std::vector<ApuTensorInfo> input_infos;
std::vector<ApuTensorInfo> output_infos;
QuantizeUtil<float, uint8_t> quantize_util_;
};
} // namespace mace
......
......@@ -46,7 +46,8 @@ class HexagonControlWrapper {
virtual bool Init() = 0;
virtual bool Finalize() = 0;
virtual bool SetupGraph(const NetDef &net_def,
const unsigned char *model_data) = 0;
const unsigned char *model_data,
const index_t model_data_size) = 0;
virtual bool ExecuteGraph(const Tensor &input_tensor,
Tensor *output_tensor) = 0;
virtual bool ExecuteGraphNew(
......
......@@ -165,7 +165,8 @@ bool HexagonDSPWrapper::Finalize() {
}
bool HexagonDSPWrapper::SetupGraph(const NetDef &net_def,
unsigned const char *model_data) {
unsigned const char *model_data,
const index_t model_data_size) {
LOG(INFO) << "Hexagon setup graph";
int64_t t0 = NowMicros();
......@@ -188,6 +189,9 @@ bool HexagonDSPWrapper::SetupGraph(const NetDef &net_def,
data = const_cast<unsigned char *>(model_data + const_tensor.offset());
data_len =
const_tensor.data_size() * GetEnumTypeSize(const_tensor.data_type());
MACE_CHECK(const_tensor.offset() + data_len <= model_data_size,
"tensor end (", const_tensor.offset() + data_len,
") should <= ", model_data_size);
}
MACE_CHECK(
hexagon_nn_append_const_node(nn_id_,
......
......@@ -35,7 +35,8 @@ class HexagonDSPWrapper : public HexagonControlWrapper {
bool Init() override;
bool Finalize() override;
bool SetupGraph(const NetDef &net_def,
const unsigned char *model_data) override;
const unsigned char *model_data,
const index_t model_data_size) override;
bool ExecuteGraph(const Tensor &input_tensor,
Tensor *output_tensor) override;
bool ExecuteGraphNew(const std::map<std::string, Tensor*> &input_tensors,
......
......@@ -138,7 +138,8 @@ bool HexagonHTAWrapper::Finalize() {
}
bool HexagonHTAWrapper::SetupGraph(const NetDef &net_def,
unsigned const char *model_data) {
unsigned const char *model_data,
const index_t model_data_size) {
LOG(INFO) << "Hexagon setup graph";
int64_t t0 = NowMicros();
......@@ -159,6 +160,9 @@ bool HexagonHTAWrapper::SetupGraph(const NetDef &net_def,
const_cast<unsigned char *>(model_data + const_tensor.offset());
const_node_data_len = const_tensor.data_size() *
GetEnumTypeSize(const_tensor.data_type());
MACE_CHECK(const_tensor.offset() + const_node_data_len <= model_data_size,
"tensor end (", const_tensor.offset() + const_node_data_len,
") should <= ", model_data_size);
}
hexagon_hta_nn_append_const_node(nn_id_,
......
......@@ -59,7 +59,8 @@ class HexagonHTAWrapper : public HexagonControlWrapper {
bool Init() override;
bool Finalize() override;
bool SetupGraph(const NetDef &net_def,
const unsigned char *model_data) override;
const unsigned char *model_data,
const index_t model_data_size) override;
bool ExecuteGraph(const Tensor &input_tensor,
Tensor *output_tensor) override;
bool ExecuteGraphNew(const std::map<std::string, Tensor*> &input_tensors,
......
......@@ -276,6 +276,7 @@ class Tensor {
inline void *raw_mutable_data() {
MACE_CHECK_NOTNULL(buffer_);
return buffer_->raw_mutable_data();
}
......
......@@ -64,6 +64,18 @@ void DequantizeTensor(Device *device,
dequantized_data);
}
index_t GetModelValidSize(const NetDef &net_def) {
index_t valid_data_size = 0;
for (auto &const_tensor : net_def.tensors()) {
valid_data_size = std::max<index_t>(
valid_data_size, const_tensor.offset() +
const_tensor.data_size()
* GetEnumTypeSize(const_tensor.data_type()));
}
return valid_data_size;
}
} // namespace
Workspace::Workspace(const OpDelegatorRegistry *registry) :
......@@ -105,23 +117,17 @@ std::vector<std::string> Workspace::Tensors() const {
return names;
}
MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
Device *device,
const unsigned char *model_data) {
MaceStatus Workspace::LoadModelTensor(const NetDef &net_def, Device *device,
const unsigned char *model_data,
const index_t model_data_size) {
MACE_LATENCY_LOGGER(1, "Load model tensors");
index_t model_data_size = 0;
for (auto &const_tensor : net_def.tensors()) {
model_data_size = std::max(
model_data_size,
static_cast<index_t>(const_tensor.offset() +
const_tensor.data_size() *
GetEnumTypeSize(const_tensor.data_type())));
}
VLOG(3) << "Model data size: " << model_data_size;
index_t valid_data_size = GetModelValidSize(net_def);
VLOG(3) << "Model valid data size: " << valid_data_size;
MACE_CHECK(valid_data_size <= model_data_size,
valid_data_size, "should be smaller than", model_data_size);
const DeviceType device_type = device->device_type();
if (model_data_size > 0) {
if (valid_data_size > 0) {
bool is_quantize_model = IsQuantizedModel(net_def);
diffused_buffer_ =
(device_type == DeviceType::CPU && HasHalfTensor(net_def)) ||
......@@ -129,7 +135,7 @@ MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
#ifdef MACE_ENABLE_OPENCL
diffused_buffer_ = diffused_buffer_ || (device_type == DeviceType::GPU &&
device->gpu_runtime()->opencl_runtime()->GetDeviceMaxMemAllocSize() <=
static_cast<uint64_t>(model_data_size));
static_cast<uint64_t>(valid_data_size));
#endif
if (diffused_buffer_) {
for (auto &const_tensor : net_def.tensors()) {
......@@ -162,15 +168,10 @@ MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
MACE_CHECK(tensor->size() == const_tensor.data_size(),
"Tensor's data_size not equal with the shape");
MACE_CHECK(static_cast<index_t>(const_tensor.offset() +
tensor->size() * GetEnumTypeSize(const_tensor.data_type())) <=
model_data_size,
"buffer offset + length (",
const_tensor.offset(),
" + ",
tensor->size() * GetEnumTypeSize(const_tensor.data_type()),
") should <= ",
model_data_size);
const index_t tensor_end = const_tensor.offset() +
tensor->size() * GetEnumTypeSize(const_tensor.data_type());
MACE_CHECK(tensor_end <= model_data_size, "tensor_end (", tensor_end,
") should <= ", model_data_size);
if (device_type == DeviceType::CPU &&
const_tensor.data_type() == DataType::DT_HALF) {
......@@ -207,14 +208,14 @@ MaceStatus Workspace::LoadModelTensor(const NetDef &net_def,
tensor_buffer_ = std::unique_ptr<Buffer>(
new Buffer(device->allocator(),
const_cast<unsigned char *>(model_data),
model_data_size));
valid_data_size));
} else {
tensor_buffer_ = std::unique_ptr<Buffer>(
new Buffer(device->allocator()));
MACE_RETURN_IF_ERROR(tensor_buffer_->Allocate(model_data_size));
MACE_RETURN_IF_ERROR(tensor_buffer_->Allocate(valid_data_size));
tensor_buffer_->Map(nullptr);
tensor_buffer_->Copy(const_cast<unsigned char *>(model_data),
0, model_data_size);
0, valid_data_size);
tensor_buffer_->UnMap();
}
for (auto &const_tensor : net_def.tensors()) {
......
......@@ -56,9 +56,9 @@ class Workspace {
std::vector<std::string> Tensors() const;
MaceStatus LoadModelTensor(const NetDef &net_def,
Device *device,
const unsigned char *model_data);
MaceStatus LoadModelTensor(const NetDef &net_def, Device *device,
const unsigned char *model_data,
const index_t model_data_size);
MaceStatus PreallocateOutputTensor(const NetDef &net_def,
const MemoryOptimizer *mem_optimizer,
......
......@@ -350,7 +350,8 @@ MaceStatus BMNet::Run(DeviceType device,
}
MaceEngine engine(config);
status = engine.Init(&net_, input_names_, output_names_, weight_.data());
status = engine.Init(&net_, input_names_, output_names_,
weight_.data(), weight_.size());
if (status != MaceStatus::MACE_SUCCESS) {
return status;
}
......
......@@ -433,7 +433,8 @@ class MaceEngine::Impl {
MaceStatus Init(const NetDef *net_def,
const std::vector<std::string> &input_nodes,
const std::vector<std::string> &output_nodes,
const unsigned char *model_data);
const unsigned char *model_data,
const int64_t model_data_size);
MaceStatus Init(const NetDef *net_def,
const std::vector<std::string> &input_nodes,
......@@ -557,7 +558,7 @@ MaceStatus MaceEngine::Impl::Init(
const NetDef *net_def,
const std::vector<std::string> &input_nodes,
const std::vector<std::string> &output_nodes,
const unsigned char *model_data) {
const unsigned char *model_data, const int64_t model_data_size) {
LOG(INFO) << "Initializing MaceEngine";
// Check avalibility
#ifdef MACE_ENABLE_OPENCL
......@@ -634,8 +635,8 @@ MaceStatus MaceEngine::Impl::Init(
MACE_CHECK(hexagon_controller_->Init(), "hexagon init error");
hexagon_controller_->SetDebugLevel(
static_cast<int>(mace::port::MinVLogLevelFromEnv()));
MACE_CHECK(hexagon_controller_->SetupGraph(*net_def, model_data),
"hexagon setup graph error");
MACE_CHECK(hexagon_controller_->SetupGraph(
*net_def, model_data, model_data_size), "hexagon setup graph error");
if (VLOG_IS_ON(2)) {
hexagon_controller_->PrintGraph();
}
......@@ -644,12 +645,12 @@ MaceStatus MaceEngine::Impl::Init(
#ifdef MACE_ENABLE_APU
if (device_type_ == APU) {
apu_controller_.reset(new ApuWrapper(device_.get()));
MACE_CHECK(apu_controller_->Init(*net_def, model_data), "apu init error");
MACE_CHECK(apu_controller_->Init(
*net_def, model_data, model_data_size), "apu init error");
} else {
#endif
MACE_RETURN_IF_ERROR(ws_->LoadModelTensor(*net_def,
device_.get(),
model_data));
MACE_RETURN_IF_ERROR(ws_->LoadModelTensor(
*net_def, device_.get(), model_data, model_data_size));
NetDef adapted_net_def;
NetDefAdapter net_def_adapter(op_registry_.get(), ws_.get());
......@@ -696,7 +697,8 @@ MaceStatus MaceEngine::Impl::Init(
MACE_RETURN_IF_ERROR(Init(
net_def, input_nodes, output_nodes,
reinterpret_cast<const unsigned char *>(model_data_->data())));
reinterpret_cast<const unsigned char *>(model_data_->data()),
model_data_->length()));
if (device_type_ == DeviceType::GPU || device_type_ == DeviceType::HEXAGON ||
device_type_ == DeviceType::HTA ||
......@@ -945,6 +947,7 @@ MaceStatus MaceEngine::Impl::Run(
<< MakeString(MapKeys(output_info_map_));
}
Tensor *output_tensor = ws_->GetTensor(output.first);
MACE_CHECK(output_tensor != nullptr);
output_tensors[output.first] = output_tensor;
}
#if defined(MACE_ENABLE_HEXAGON) || defined(MACE_ENABLE_HTA)
......@@ -988,8 +991,10 @@ MaceEngine::~MaceEngine() = default;
MaceStatus MaceEngine::Init(const NetDef *net_def,
const std::vector<std::string> &input_nodes,
const std::vector<std::string> &output_nodes,
const unsigned char *model_data) {
return impl_->Init(net_def, input_nodes, output_nodes, model_data);
const unsigned char *model_data,
int64_t model_data_size) {
return impl_->Init(net_def, input_nodes, output_nodes,
model_data, model_data_size);
}
MaceStatus MaceEngine::Init(const NetDef *net_def,
......@@ -1032,7 +1037,8 @@ MaceStatus CreateMaceEngineFromProto(
engine->reset(new mace::MaceEngine(config));
MaceStatus status = (*engine)->Init(
net_def.get(), input_nodes, output_nodes, model_weights_data);
net_def.get(), input_nodes, output_nodes,
model_weights_data, model_weights_data_size);
return status;
}
......
......@@ -31,7 +31,9 @@ void CalcPaddingAndOutputSize(const index_t *input_shape,
index_t *output_shape,
int *padding_size) {
MACE_CHECK(dilations[0] > 0 && dilations[1] > 0,
"Invalid dilations, must >= 1");
"Invalid dilations, must > 0");
MACE_CHECK(strides[0] > 0 && strides[1] > 0,
"Invalid strides, must > 0");
MACE_CHECK((dilations[0] == 1 || strides[0] == 1) &&
(dilations[1] == 1 || strides[1] == 1),
"If dilations > 1, strides should be 1");
......@@ -145,7 +147,9 @@ void CalcOutputSize(const index_t *input_shape,
const RoundType round_type,
index_t *output_shape) {
MACE_CHECK(dilations[0] > 0 && dilations[1] > 0,
"Invalid dilations, must >= 1");
"Invalid dilations, must > 0");
MACE_CHECK(strides[0] > 0 && strides[1] > 0,
"Invalid strides, must > 0");
MACE_CHECK((dilations[0] == 1 || strides[0] == 1) &&
(dilations[1] == 1 || strides[1] == 1),
"If dilations > 1, strides should be 1");
......
......@@ -94,6 +94,7 @@ class PoolingOp<DeviceType::CPU, T> : public PoolingOpBase {
Tensor::MappingGuard input_guard(input_tensor);
Tensor::MappingGuard output_guard(output_tensor);
const T *input = input_tensor->data<T>();
MACE_CHECK(output_tensor->dtype() == DataTypeToEnum<T>::value);
T *output = output_tensor->mutable_data<T>();
const index_t *input_shape = input_tensor->shape().data();
int pad_hw[2] = {paddings[0] / 2, paddings[1] / 2};
......@@ -299,6 +300,7 @@ class PoolingOp<DeviceType::CPU, uint8_t> : public PoolingOpBase {
Tensor::MappingGuard input_guard(input_tensor);
Tensor::MappingGuard output_guard(output_tensor);
const uint8_t *input = input_tensor->data<uint8_t>();
MACE_CHECK(output_tensor->dtype() == DataTypeToEnum<uint8_t>::value);
uint8_t *output = output_tensor->mutable_data<uint8_t>();
int pad_hw[2] = {paddings[0] / 2, paddings[1] / 2};
......
......@@ -39,9 +39,9 @@ MaceStatus Conv2d<T>::Compute(const OpContext *context,
const Tensor *filter,
Tensor *output) {
MACE_UNUSED(context);
const std::vector<index_t> in_shape = input->shape();
const std::vector<index_t> filter_shape = filter->shape();
MACE_CHECK(in_shape[1] == filter_shape[1]);
std::vector<index_t> out_shape(4);
std::vector<int> paddings(2);
......
......@@ -77,6 +77,7 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
{% if embed_model_data %}
(void)model_data_file;
const unsigned char * model_data;
const int64_t model_size;
{% endif %}
MaceStatus status = MaceStatus::MACE_SUCCESS;
switch (model_name_map[model_name]) {
......@@ -86,8 +87,9 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
engine->reset(new mace::MaceEngine(config));
{% if embed_model_data %}
model_data = mace::{{model_tags[i]}}::LoadModelData();
model_size = mace::{{model_tags[i]}}::GetModelSize();
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_data);
model_data, model_size);
{% else %}
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_data_file);
......@@ -116,10 +118,9 @@ MaceStatus CreateMaceEngineFromCode(
std::shared_ptr<NetDef> net_def;
{% if embed_model_data %}
const unsigned char * model_data;
const int64_t model_size;
(void)model_weights_data;
{% endif %}
// TODO(yejianwu) Add buffer range checking
(void)model_weights_data_size;
MaceStatus status = MaceStatus::MACE_SUCCESS;
switch (model_name_map[model_name]) {
......@@ -129,11 +130,12 @@ MaceStatus CreateMaceEngineFromCode(
engine->reset(new mace::MaceEngine(config));
{% if embed_model_data %}
model_data = mace::{{model_tags[i]}}::LoadModelData();
model_size = mace::{{model_tags[i]}}::GetModelSize();
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_data);
model_data, model_size);
{% else %}
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_weights_data);
model_weights_data, model_weights_data_size);
{% endif %}
break;
{% endfor %}
......
......@@ -36,7 +36,7 @@ TEST(MaceAPIExceptionTest, WrongInputTest) {
}
MaceEngine engine(config);
ASSERT_DEATH(engine.Init(net_def.get(), {"input"}, output_names, nullptr),
ASSERT_DEATH(engine.Init(net_def.get(), {"input"}, output_names, nullptr, 0),
"");
}
......
......@@ -68,8 +68,10 @@ void MaceRunFunc(const int in_out_size) {
config.SetGPUContext(mace::ops::test::OpTestContext::Get()->gpu_context());
MaceEngine engine(config);
MaceStatus status = engine.Init(net_def.get(), input_names, output_names,
reinterpret_cast<unsigned char *>(data.data()));
MaceStatus status = engine.Init(
net_def.get(), input_names, output_names,
reinterpret_cast<unsigned char *>(data.data()),
data.size() * sizeof(half));
EXPECT_EQ(status, MaceStatus::MACE_SUCCESS);
std::map<std::string, mace::MaceTensor> inputs;
......
......@@ -65,7 +65,7 @@ void MaceRun(const int in_out_size,
MaceEngine engine(config);
MaceStatus status = engine.Init(net_def.get(), input_names, output_names,
reinterpret_cast<unsigned char *>(data.data()));
reinterpret_cast<unsigned char *>(data.data()), data.size() * sizeof(T));
EXPECT_EQ(status, MaceStatus::MACE_SUCCESS);
std::map<std::string, mace::MaceTensor> inputs;
......
// Copyright 2018 The MACE Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mace/ops/memory_optimizer_for_test.h"
namespace mace {
void MemoryOptimizerForTest::Optimize(
const OperatorDef *op_def,
const std::unordered_map<std::string, MemoryType> *mem_types) {
if (op_def->output_size() == op_def->output_shape_size()) {
MemoryOptimizer::Optimize(op_def, mem_types);
}
}
} // namespace mace
// Copyright 2020 The MACE Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_
#define MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_
#include <string>
#include <unordered_map>
#include "mace/core/memory_optimizer.h"
namespace mace {
class MemoryOptimizerForTest : public MemoryOptimizer {
public:
MemoryOptimizerForTest() {}
~MemoryOptimizerForTest() {}
void Optimize(
const OperatorDef *op_def,
const std::unordered_map<std::string, MemoryType> *mem_types) override;
};
} // namespace mace
#endif // MACE_OPS_MEMORY_OPTIMIZER_FOR_TEST_H_
......@@ -13,9 +13,11 @@
// limitations under the License.
#include "mace/ops/ops_test_util.h"
#include "mace/core/memory_optimizer.h"
#include "mace/utils/memory.h"
#include "mace/core/net_def_adapter.h"
#include "mace/ops/memory_optimizer_for_test.h"
#include "mace/utils/memory.h"
namespace mace {
namespace ops {
......@@ -228,7 +230,7 @@ bool OpsTestNet::Setup(mace::DeviceType device) {
OpTestContext::Get()->GetDevice(device),
&adapted_net_def);
MemoryOptimizer mem_optimizer;
MemoryOptimizerForTest mem_optimizer;
net_ = make_unique<SerialNet>(
op_registry_.get(),
&adapted_net_def,
......@@ -284,7 +286,7 @@ MaceStatus OpsTestNet::RunNet(const mace::NetDef &net_def,
net_def_adapter.AdaptNetDef(&net_def,
OpTestContext::Get()->GetDevice(device),
&adapted_net_def);
MemoryOptimizer mem_optimizer;
MemoryOptimizerForTest mem_optimizer;
net_ = make_unique<SerialNet>(
op_registry_.get(),
&adapted_net_def,
......
......@@ -77,6 +77,7 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
{% if embed_model_data %}
(void)model_data_file;
const unsigned char * model_data;
const int64_t model_size;
{% endif %}
MaceStatus status = MaceStatus::MACE_SUCCESS;
switch (model_name_map[model_name]) {
......@@ -86,8 +87,9 @@ __attribute__((deprecated)) MaceStatus CreateMaceEngineFromCode(
engine->reset(new mace::MaceEngine(config));
{% if embed_model_data %}
model_data = mace::{{model_tags[i]}}::LoadModelData();
model_size = mace::{{model_tags[i]}}::GetModelSize();
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_data);
model_data, model_size);
{% else %}
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_data_file);
......@@ -116,10 +118,9 @@ MACE_API MaceStatus CreateMaceEngineFromCode(
std::shared_ptr<NetDef> net_def;
{% if embed_model_data %}
const unsigned char * model_data;
const int64_t model_size;
(void)model_weights_data;
{% endif %}
// TODO(yejianwu) Add buffer range checking
(void)model_weights_data_size;
MaceStatus status = MaceStatus::MACE_SUCCESS;
switch (model_name_map[model_name]) {
......@@ -129,11 +130,12 @@ MACE_API MaceStatus CreateMaceEngineFromCode(
engine->reset(new mace::MaceEngine(config));
{% if embed_model_data %}
model_data = mace::{{model_tags[i]}}::LoadModelData();
model_size = mace::{{model_tags[i]}}::GetModelSize();
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_data);
model_data, model_size);
{% else %}
status = (*engine)->Init(net_def.get(), input_nodes, output_nodes,
model_weights_data);
model_weights_data, model_weights_data_size);
{% endif %}
break;
{% endfor %}
......
......@@ -27,6 +27,10 @@ MACE_API const unsigned char *LoadModelData() {
return model_data;
}
MACE_API int64_t GetModelSize() {
return {{ model_data_size }};
}
} // namespace {{tag}}
} // namespace mace
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册