wehub-resource-sync
367 lines
13 KiB
C++
367 lines
13 KiB
C++
//
|
|
// ChannelPruneConvert.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2023/05/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "CommonUtils.hpp"
|
|
#include "MNN/expr/ExprCreator.hpp"
|
|
#include <vector>
|
|
#include <map>
|
|
#include <set>
|
|
#include <algorithm>
|
|
|
|
using namespace MNN;
|
|
using namespace std;
|
|
|
|
// TODO: add more unsafe ops
|
|
static std::vector<MNN::OpType> unSafeOpTypes = {
|
|
OpType_BroadcastTo, OpType_BatchToSpaceND, OpType_Concat, OpType_LSTM, OpType_LSTMBlockCell, OpType_Reshape, OpType_Resize,
|
|
OpType_RNN, OpType_RNNSequenceGRU, OpType_ScatterNd, OpType_Slice, OpType_SliceTf, OpType_SpaceToBatchND, OpType_Raster,
|
|
};
|
|
|
|
struct TensorMaskInfo {
|
|
std::vector<int> mask; // per-channel 1 or 0
|
|
std::string oriConvName;
|
|
};
|
|
|
|
std::vector<MNN::OpT*> findUserOps(int outputIndex, std::unique_ptr<MNN::NetT>& netT, SubGraphProtoT* subgraph) {
|
|
std::vector<MNN::OpT*> userOps;
|
|
if (subgraph) {
|
|
for (auto& subOp : subgraph->nodes) {
|
|
for (int inputIndex : subOp->inputIndexes) {
|
|
if (inputIndex == outputIndex) {
|
|
userOps.push_back(subOp.get());
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
for (auto& netOp : netT->oplists) {
|
|
for (int inputIndex : netOp->inputIndexes) {
|
|
if (inputIndex == outputIndex) {
|
|
userOps.push_back(netOp.get());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return userOps;
|
|
}
|
|
|
|
// do the actual channel prune on weights and bias
|
|
void channelPrune(std::unique_ptr<MNN::OpT>& op, std::unique_ptr<MNN::NetT>& netT, SubGraphProtoT* subgraph, std::map<std::string, TensorMaskInfo>& tensorMaskInfo) {
|
|
auto opType = op->type;
|
|
if (opType != OpType_Convolution && opType != OpType_ConvolutionDepthwise && opType != OpType_Deconvolution && opType != OpType_DeconvolutionDepthwise && opType != OpType_BatchNorm) {
|
|
return;
|
|
}
|
|
if (op->inputIndexes.size() != 1) {
|
|
return;
|
|
}
|
|
|
|
int inputIndex = op->inputIndexes[0];
|
|
int outputIndex = op->outputIndexes[0];
|
|
std::string inputTensorName = subgraph ? subgraph->tensors[inputIndex] : netT->tensorName[inputIndex];
|
|
std::string outputTensorName = subgraph ? subgraph->tensors[outputIndex] : netT->tensorName[outputIndex];
|
|
|
|
std::vector<int> inputMask = tensorMaskInfo[inputTensorName].mask;
|
|
int inputMaskSum = 0;
|
|
for (int i = 0; i < inputMask.size(); i++) {
|
|
inputMaskSum += inputMask[i];
|
|
}
|
|
|
|
if (opType == OpType_BatchNorm) {
|
|
if (!(inputMaskSum < inputMask.size())) {
|
|
return;
|
|
}
|
|
|
|
auto bnParams = op->main.AsBatchNorm();
|
|
auto slopFloat = bnParams->slopeData;
|
|
auto biasFloat = bnParams->biasData;
|
|
auto meanFloat = bnParams->meanData;
|
|
auto varianceFloat = bnParams->varData;
|
|
|
|
bnParams->slopeData.clear();
|
|
bnParams->biasData.clear();
|
|
bnParams->meanData.clear();
|
|
bnParams->varData.clear();
|
|
|
|
for (int i = 0; i < varianceFloat.size(); i++) {
|
|
if (inputMask[i] == 1) {
|
|
bnParams->slopeData.push_back(slopFloat[i]);
|
|
bnParams->biasData.push_back(biasFloat[i]);
|
|
bnParams->meanData.push_back(meanFloat[i]);
|
|
bnParams->varData.push_back(varianceFloat[i]);
|
|
}
|
|
}
|
|
bnParams->channels = inputMaskSum;
|
|
|
|
return;
|
|
}
|
|
|
|
auto convParams = op->main.AsConvolution2D();
|
|
auto weightFloat = convParams->weight;
|
|
auto biasFloat = convParams->bias;
|
|
auto& common = convParams->common;
|
|
|
|
int ko = common->outputCount;
|
|
int ki = common->inputCount / common->group;
|
|
int kh = common->kernelY;
|
|
int kw = common->kernelX;
|
|
|
|
std::vector<int> opMask;
|
|
for (auto info : tensorMaskInfo) {
|
|
if (op->name == info.second.oriConvName) {
|
|
opMask = info.second.mask;
|
|
break;
|
|
}
|
|
}
|
|
|
|
int opMaskSum = 0;
|
|
for (int i = 0; i < opMask.size(); i++) {
|
|
opMaskSum += opMask[i];
|
|
}
|
|
|
|
if (opMaskSum < opMask.size()) {
|
|
convParams->weight.clear();
|
|
convParams->bias.clear();
|
|
|
|
for (int i = 0; i < ko; i++) {
|
|
int offset = i * ki * kh * kw;
|
|
if (opMask[i] == 1) {
|
|
for (int j = 0; j < ki * kh * kw; j++) {
|
|
convParams->weight.emplace_back(weightFloat[offset + j]);
|
|
}
|
|
convParams->bias.emplace_back(biasFloat[i]);
|
|
}
|
|
}
|
|
common->outputCount = opMaskSum;
|
|
}
|
|
|
|
if (inputMaskSum < inputMask.size()) {
|
|
auto weightFloat = convParams->weight;
|
|
convParams->weight.clear();
|
|
|
|
int ko = common->outputCount;
|
|
int ki = common->inputCount / common->group;
|
|
int kh = common->kernelY;
|
|
int kw = common->kernelX;
|
|
|
|
for (int i = 0; i < ko; i++) {
|
|
for (int j = 0; j < ki; j++) {
|
|
int offset = i * ki * kh * kw + j * kh * kw;
|
|
if (inputMask[j] == 1) {
|
|
for (int k = 0; k < kh * kw; k++) {
|
|
convParams->weight.emplace_back(weightFloat[offset + k]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
common->inputCount = inputMaskSum;
|
|
|
|
// we will not do prune for depthwise, its channel pruning only depends on its input tensor's pruning
|
|
if (opType == OpType_ConvolutionDepthwise || opType == OpType_DeconvolutionDepthwise) {
|
|
common->outputCount = inputMaskSum;
|
|
}
|
|
}
|
|
}
|
|
|
|
// propagate and analyze prune mask info in model
|
|
void analyzePruneInfo(std::unique_ptr<MNN::OpT>& op, std::unique_ptr<MNN::NetT>& netT, SubGraphProtoT* subgraph, std::map<std::string, TensorMaskInfo>& tensorMaskInfo, std::set<std::string>& notSafeConvNames) {
|
|
auto opType = op->type;
|
|
auto inputIndices = op->inputIndexes;
|
|
if (inputIndices.size() == 0) {
|
|
return;
|
|
}
|
|
auto outputIndices = op->outputIndexes;
|
|
std::vector<std::string> inputTensorNames;
|
|
for (int i = 0; i < inputIndices.size(); i++) {
|
|
inputTensorNames.push_back(subgraph ? subgraph->tensors[inputIndices[i]] : netT->tensorName[inputIndices[i]]);
|
|
}
|
|
std::vector<std::string> outputTensorNames;
|
|
for (int i = 0; i < outputIndices.size(); i++) {
|
|
outputTensorNames.push_back(subgraph ? subgraph->tensors[outputIndices[i]] : netT->tensorName[outputIndices[i]]);
|
|
}
|
|
|
|
if (opType == OpType_Convolution || opType == OpType_Deconvolution) {
|
|
if (inputIndices.size() == 1) {
|
|
auto convParams = op->main.AsConvolution2D();
|
|
auto weightFloat = convParams->weight;
|
|
auto biasFloat = convParams->bias;
|
|
auto& common = convParams->common;
|
|
|
|
const int ko = common->outputCount;
|
|
const int ki = common->inputCount / common->group;
|
|
const int kh = common->kernelY;
|
|
const int kw = common->kernelX;
|
|
|
|
MNN::Express::VARP weightVar = MNN::Express::_Const(weightFloat.data(), {ko, ki, kh, kw}, MNN::Express::NCHW);
|
|
|
|
MNN::Express::VARP weightMask = MNN::Express::_Greater(MNN::Express::_ReduceSum(MNN::Express::_Abs(weightVar), {1, 2, 3}), MNN::Express::_Scalar<float>(1e-6));
|
|
MNN::Express::VARP maskSum = MNN::Express::_ReduceSum(weightMask);
|
|
auto maskInfo = weightMask->getInfo();
|
|
auto maskPtr = weightMask->readMap<int>();
|
|
|
|
if (maskSum->readMap<int>()[0] == maskInfo->size) {
|
|
return;
|
|
}
|
|
|
|
// conv has pruned, propagate its mask down
|
|
tensorMaskInfo[outputTensorNames[0]].oriConvName = op->name;
|
|
for (int i = 0; i < maskInfo->size; i++) {
|
|
tensorMaskInfo[outputTensorNames[0]].mask.push_back(maskPtr[i]);
|
|
}
|
|
}
|
|
|
|
return;
|
|
}
|
|
|
|
std::vector<MNN::OpType>::iterator iter;
|
|
iter = std::find(unSafeOpTypes.begin(), unSafeOpTypes.end(), opType);
|
|
// not safe op and num_outputs > 1 op are not safe
|
|
if ((iter != unSafeOpTypes.end()) || (outputTensorNames.size() > 1)) {
|
|
for (auto name : inputTensorNames) {
|
|
if (!tensorMaskInfo[name].oriConvName.empty()) {
|
|
// so that input tensor mask's oriConv op is not safe
|
|
notSafeConvNames.insert(tensorMaskInfo[name].oriConvName);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
// when a mask is propagated to the output, its oriConv ops are not safe
|
|
std::vector<MNN::OpT*> userOps = findUserOps(outputIndices[0], netT, subgraph);
|
|
if (userOps.size() == 0) {
|
|
for (auto name : inputTensorNames) {
|
|
if (!tensorMaskInfo[name].oriConvName.empty()) {
|
|
notSafeConvNames.insert(tensorMaskInfo[name].oriConvName);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
// if the op has more than one input (including const input)
|
|
// we need its input tensor's masks are all from one oriConv op
|
|
if (inputIndices.size() > 1) {
|
|
std::string oriConvName;
|
|
std::string oriTensorName;
|
|
for (auto name : inputTensorNames) {
|
|
if (!tensorMaskInfo[name].oriConvName.empty()) {
|
|
oriConvName = tensorMaskInfo[name].oriConvName;
|
|
oriTensorName = name;
|
|
}
|
|
}
|
|
if (oriConvName.empty()) {
|
|
return;
|
|
}
|
|
|
|
// oriConvName is not empty
|
|
bool unsafe = false;
|
|
for (auto name : inputTensorNames) {
|
|
auto tOriName = tensorMaskInfo[name].oriConvName;
|
|
if ((tOriName != oriConvName) && (!tOriName.empty())) {
|
|
unsafe = true;
|
|
}
|
|
}
|
|
|
|
// if unsafe, all its input tensor mask's oriConvs are not safe
|
|
if (unsafe) {
|
|
for (auto name : inputTensorNames) {
|
|
auto tOriName = tensorMaskInfo[name].oriConvName;
|
|
if (!tOriName.empty()) {
|
|
notSafeConvNames.insert(tOriName);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
// if safe, propagate mask down
|
|
tensorMaskInfo[outputTensorNames[0]].oriConvName = oriConvName;
|
|
tensorMaskInfo[outputTensorNames[0]].mask = tensorMaskInfo[oriTensorName].mask;
|
|
return;
|
|
}
|
|
|
|
// for 1 input and 1 output safe op, propagate mask down
|
|
tensorMaskInfo[outputTensorNames[0]].oriConvName = tensorMaskInfo[inputTensorNames[0]].oriConvName;
|
|
tensorMaskInfo[outputTensorNames[0]].mask = tensorMaskInfo[inputTensorNames[0]].mask;
|
|
}
|
|
|
|
void channelPruneConvert(std::unique_ptr<MNN::NetT>& netT, MNN::Compression::Pipeline proto) {
|
|
bool filterPruned = false;
|
|
for (const auto& algo : proto.algo()) {
|
|
if (algo.type() == Compression::CompressionAlgo::PRUNE) {
|
|
auto prune_type = algo.prune_params().type();
|
|
auto prune_algo_type = MNN::SparseAlgo(prune_type);
|
|
if (prune_type == Compression::PruneParams_PruneType_FILTER) {
|
|
filterPruned = true;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!filterPruned) {
|
|
return;
|
|
}
|
|
|
|
std::map<std::string, TensorMaskInfo> netMaskInfo;
|
|
for (auto tensorName : netT->tensorName) {
|
|
netMaskInfo[tensorName] = TensorMaskInfo();
|
|
}
|
|
|
|
std::set<std::string> notSafeConvNames;
|
|
for (auto& op : netT->oplists) {
|
|
analyzePruneInfo(op, netT, nullptr, netMaskInfo, notSafeConvNames);
|
|
}
|
|
|
|
std::set<std::string>::iterator iter;
|
|
if (!notSafeConvNames.empty()) {
|
|
for (auto& info : netMaskInfo) {
|
|
iter = std::find(notSafeConvNames.begin(), notSafeConvNames.end(), info.second.oriConvName);
|
|
if (iter != notSafeConvNames.end()) {
|
|
for (int i = 0; i < info.second.mask.size(); i++) {
|
|
if (info.second.mask[i] == 0) {
|
|
info.second.mask[i] = 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for (auto& op : netT->oplists) {
|
|
channelPrune(op, netT, nullptr, netMaskInfo);
|
|
}
|
|
|
|
|
|
for (auto& subgraph : netT->subgraphs) {
|
|
std::map<std::string, TensorMaskInfo> subgraphMaskInfo;
|
|
for (auto tensorName : subgraph->tensors) {
|
|
subgraphMaskInfo[tensorName] = TensorMaskInfo();
|
|
}
|
|
|
|
std::set<std::string> notSafeConvNames;
|
|
for (auto& op : subgraph->nodes) {
|
|
analyzePruneInfo(op, netT, subgraph.get(), subgraphMaskInfo, notSafeConvNames);
|
|
}
|
|
|
|
std::set<std::string>::iterator iter;
|
|
if (!notSafeConvNames.empty()) {
|
|
for (auto& info : subgraphMaskInfo) {
|
|
iter = std::find(notSafeConvNames.begin(), notSafeConvNames.end(), info.second.oriConvName);
|
|
if (iter != notSafeConvNames.end()) {
|
|
for (int i = 0; i < info.second.mask.size(); i++) {
|
|
if (info.second.mask[i] == 0) {
|
|
info.second.mask[i] = 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for (auto& op : subgraph->nodes) {
|
|
channelPrune(op, netT, subgraph.get(), subgraphMaskInfo);
|
|
}
|
|
}
|
|
}
|