Files
2026-07-13 13:33:03 +08:00

143 lines
4.1 KiB
C++

//
// MobilenetV1.cpp
// MNN
//
// Created by MNN on 2020/01/08.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "MobilenetV1.hpp"
#include "Initializer.hpp"
using namespace MNN::Express;
namespace MNN {
namespace Train {
namespace Model {
class _ConvBlock : public Module {
public:
_ConvBlock(std::vector<int> inputOutputChannels, int stride);
virtual std::vector<Express::VARP> onForward(const std::vector<Express::VARP> &inputs) override;
std::shared_ptr<Module> conv3x3;
std::shared_ptr<Module> bn1;
std::shared_ptr<Module> conv1x1;
std::shared_ptr<Module> bn2;
};
std::shared_ptr<Module> ConvBlock(std::vector<int> inputOutputChannels, int stride) {
return std::shared_ptr<Module>(new _ConvBlock(inputOutputChannels, stride));
}
_ConvBlock::_ConvBlock(std::vector<int> inputOutputChannels, int stride) {
int inputChannels = inputOutputChannels[0], outputChannels = inputOutputChannels[1];
NN::ConvOption convOption;
convOption.kernelSize = {3, 3};
convOption.channel = {inputChannels, inputChannels};
convOption.padMode = Express::SAME;
convOption.stride = {stride, stride};
convOption.depthwise = true;
conv3x3.reset(NN::Conv(convOption, false, std::shared_ptr<Initializer>(Initializer::MSRA())));
bn1.reset(NN::BatchNorm(inputChannels));
convOption.reset();
convOption.kernelSize = {1, 1};
convOption.channel = {inputChannels, outputChannels};
convOption.padMode = Express::SAME;
convOption.stride = {1, 1};
convOption.depthwise = false;
conv1x1.reset(NN::Conv(convOption, false, std::shared_ptr<Initializer>(Initializer::MSRA())));
bn2.reset(NN::BatchNorm(outputChannels));
registerModel({conv3x3, bn1, conv1x1, bn2});
}
std::vector<Express::VARP> _ConvBlock::onForward(const std::vector<Express::VARP> &inputs) {
using namespace Express;
VARP x = inputs[0];
x = conv3x3->forward(x);
x = bn1->forward(x);
x = _Relu6(x);
x = conv1x1->forward(x);
x = bn2->forward(x);
x = _Relu6(x);
return {x};
}
MobilenetV1::MobilenetV1(int numClasses, float widthMult, int divisor) {
NN::ConvOption convOption;
convOption.kernelSize = {3, 3};
int outputChannels = makeDivisible(32 * widthMult, divisor);
convOption.channel = {3, outputChannels};
convOption.padMode = Express::SAME;
convOption.stride = {2, 2};
conv1.reset(NN::Conv(convOption, false, std::shared_ptr<Initializer>(Initializer::MSRA())));
bn1.reset(NN::BatchNorm(outputChannels));
std::vector<std::vector<int> > convSettings;
convSettings.push_back({64, 1});
convSettings.push_back({128, 2});
convSettings.push_back({256, 2});
convSettings.push_back({512, 6});
convSettings.push_back({1024, 2});
int inputChannels = outputChannels;
for (int i = 0; i < convSettings.size(); i++) {
auto setting = convSettings[i];
outputChannels = setting[0];
int times = setting[1];
outputChannels = makeDivisible(outputChannels * widthMult, divisor);
for (int j = 0; j < times; j++) {
int stride = 1;
if (times > 1 && j == 0) {
stride = 2;
}
convBlocks.emplace_back(ConvBlock({inputChannels, outputChannels}, stride));
inputChannels = outputChannels;
}
}
dropout.reset(NN::Dropout(0.1));
fc.reset(NN::Linear(1024, numClasses, true, std::shared_ptr<Initializer>(Initializer::MSRA())));
registerModel({conv1, bn1, dropout, fc});
registerModel(convBlocks);
}
std::vector<Express::VARP> MobilenetV1::onForward(const std::vector<Express::VARP> &inputs) {
using namespace Express;
VARP x = inputs[0];
x = conv1->forward(x);
x = bn1->forward(x);
x = _Relu6(x);
for (int i = 0; i < convBlocks.size(); i++) {
x = convBlocks[i]->forward(x);
}
// global avg pooling
x = _AvePool(x, {-1, -1});
x = _Convert(x, NCHW);
x = _Reshape(x, {0, -1});
x = dropout->forward(x);
x = fc->forward(x);
x = _Softmax(x, 1);
return {x};
}
} // namespace Model
} // namespace Train
} // namespace MNN