wehub-resource-sync
68 lines
2.6 KiB
C++
68 lines
2.6 KiB
C++
//
|
|
// ShuffleNetExpr.cpp
|
|
// MNN
|
|
// Reference paper: https://arxiv.org/pdf/1707.01083.pdf
|
|
//
|
|
// Created by MNN on 2019/06/25.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <map>
|
|
#include "ShuffleNetExpr.hpp"
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
|
|
using namespace MNN::Express;
|
|
|
|
// bottleNeckChannel = outputChannel / narrowRatio
|
|
static VARP shuffleUnit(VARP x, int inputChannel, int outputChannel,
|
|
int group, int stride, int narrowRatio) {
|
|
int bottleNeckChannel = outputChannel / narrowRatio;
|
|
int branchChannel = outputChannel;
|
|
if (stride != 1) {
|
|
branchChannel = outputChannel - inputChannel;
|
|
}
|
|
auto y = _Conv(0.0f, 0.0f, x, {inputChannel, bottleNeckChannel}, {1, 1}, VALID, {1, 1}, {1, 1}, group); // Group Conv
|
|
y = _ChannelShuffle(y, group);
|
|
y = _Conv(0.0f, 0.0f, y, {bottleNeckChannel, bottleNeckChannel}, {3, 3}, SAME, {stride, stride}, {1, 1}, bottleNeckChannel); // Depthwise Conv
|
|
y = _Conv(0.0f, 0.0f, y, {bottleNeckChannel, branchChannel}, {1, 1}, VALID, {1, 1}, {1, 1}, group); // Group Conv
|
|
if (stride != 1) {
|
|
x = _AvePool(x, {3, 3}, {2, 2}, SAME);
|
|
y = _Concat({x, y}, 1); // concat on channel axis (NCHW)
|
|
} else {
|
|
y = _Add(x, y);
|
|
}
|
|
return y;
|
|
}
|
|
|
|
static VARP shuffleBlock(VARP x, int inputChannel, int outputChannel,
|
|
int group, int stride, int narrowRatio, int number) {
|
|
x = shuffleUnit(x, inputChannel, outputChannel, group, stride, narrowRatio);
|
|
for (int i = 1; i < number; ++i) {
|
|
x = shuffleUnit(x, outputChannel, outputChannel, group, 1, narrowRatio);
|
|
}
|
|
return x;
|
|
}
|
|
|
|
VARP shuffleNetExpr(int group, int numClass) {
|
|
std::map<int, std::vector<int>> groupMap = {
|
|
{1, {144, 288, 576}},
|
|
{2, {200, 400, 800}},
|
|
{3, {240, 480, 960}},
|
|
{4, {272, 544, 1088}},
|
|
{8, {384, 768, 1536}}
|
|
};
|
|
int outputChannelStage2 = groupMap[group][0];
|
|
int outputChannelStage3 = groupMap[group][1];
|
|
int outputChannelStage4 = groupMap[group][2];
|
|
// Net Construction
|
|
auto x = _Input({1, 3, 224, 224}, NC4HW4);
|
|
x = _Conv(0.0f, 0.0f, x, {3, 24}, {3, 3}, SAME, {2, 2}, {1, 1}, 1);
|
|
x = _MaxPool(x, {3, 3}, {2, 2}, SAME);
|
|
x = shuffleBlock(x, 24, outputChannelStage2, group, 2, 4, 4);
|
|
x = shuffleBlock(x, outputChannelStage2, outputChannelStage3, group, 2, 4, 8);
|
|
x = shuffleBlock(x, outputChannelStage3, outputChannelStage4, group, 2, 4, 4);
|
|
x = _AvePool(x, {7, 7}, {1, 1}, VALID);
|
|
x = _Conv(0.0f, 0.0f, x, {outputChannelStage4, numClass}, {1, 1}, VALID, {1, 1}, {1, 1}, 1); // replace FC with Conv1x1
|
|
return x;
|
|
}
|