wehub-resource-sync
67 lines
1.8 KiB
C++
67 lines
1.8 KiB
C++
//
|
|
// TfUtils.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/01/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <stdio.h>
|
|
#include <algorithm>
|
|
#include <fstream>
|
|
#include <limits>
|
|
#include <set>
|
|
|
|
#include "TfUtils.hpp"
|
|
#include "logkit.h"
|
|
|
|
bool tf_read_proto_from_binary(const char* filepath, google::protobuf::Message* message) {
|
|
std::ifstream fs(filepath, std::ifstream::in | std::ifstream::binary);
|
|
if (!fs.is_open()) {
|
|
fprintf(stderr, "open failed %s\n", filepath);
|
|
return false;
|
|
}
|
|
|
|
google::protobuf::io::IstreamInputStream input(&fs);
|
|
google::protobuf::io::CodedInputStream codedstr(&input);
|
|
#if GOOGLE_PROTOBUF_VERSION >= 3011000
|
|
codedstr.SetTotalBytesLimit(INT_MAX);
|
|
#else
|
|
codedstr.SetTotalBytesLimit(INT_MAX, INT_MAX/2);
|
|
#endif
|
|
bool success = message->ParseFromCodedStream(&codedstr);
|
|
|
|
fs.close();
|
|
|
|
return success;
|
|
}
|
|
|
|
bool find_attr_value(const tensorflow::NodeDef* node, const char* key, tensorflow::AttrValue& value) {
|
|
const google::protobuf::Map<std::string, tensorflow::AttrValue>& attr = node->attr();
|
|
|
|
const google::protobuf::Map<std::string, tensorflow::AttrValue>::const_iterator it = attr.find(key);
|
|
if (it != attr.end()) {
|
|
value = it->second;
|
|
return true;
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
bool convertDataFormat(const float* src, float* dst, int planeNumber, int CI, int CO) {
|
|
// H W CI CO --> CO CI H W
|
|
assert(planeNumber > 0);
|
|
assert(CI > 0);
|
|
assert(CO > 0);
|
|
assert(src != nullptr);
|
|
for (int coi = 0; coi < CO; coi++) {
|
|
for (int cii = 0; cii < CI; cii++) {
|
|
for (int i = 0; i < planeNumber; ++i) {
|
|
dst[(coi * CI + cii) * planeNumber + i] = src[(i * CI + cii) * CO + coi];
|
|
}
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|