Files
alibaba--mnn/tools/converter/source/optimizer/merge/Convolution3DTurn2D.cpp
T
2026-07-13 13:33:03 +08:00

267 lines
12 KiB
C++

//
// ConvBiasAdd.cpp
// MNNConverter
//
// Created by MNN on 2019/09/16.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "../TemplateMerge.hpp"
#include <MNN/expr/ExprCreator.hpp>
#include "MNN_generated.h"
namespace MNN {
namespace Express {
static EXPRP _transformConv3DWithConv2D(EXPRP expr) {
std::unique_ptr<MNN::OpT> originOp(expr->get()->UnPack());
auto common = originOp->main.AsConvolution3D()->common.get();
auto weightPtr = originOp->main.AsConvolution3D()->weight.data();
auto biasDataPtr = originOp->main.AsConvolution3D()->bias.data();
auto input = expr->inputs()[0];
// Im2Col
auto one = _Unsqueeze(_Scalar<int32_t>(1), {0});
auto negone = _Unsqueeze(_Scalar<int>(-1), {0});
auto sx = _Shape(input, true);
auto kernelD = common->kernels[0];
auto kernelH = common->kernels[1];
auto kernelW = common->kernels[2];
auto inputChannel = common->inputCount;
auto outputChannel = common->outputCount;
auto kdv = _Unsqueeze(_Scalar<int>(kernelD), {0});
auto w = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(4), {0}), one);
auto h = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(3), {0}), one);
auto d = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(2), {0}), one);
auto ic = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(1), {0}), one);
auto b = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(0), {0}), one);
auto im2colinput = _Reshape(input, _Concat({one, b*ic, d, w*h}, 0));
auto im2coloutput = _Im2Col(im2colinput, {1, kernelD}, {1, common->dilates[0]}, {common->pads[0], 0, common->pads[3], 0}, {1, common->strides[0]});
// Reshape and Unpack
std::vector<VARP> convInputs;
{
// use -1 to compute od
auto value = _Reshape(im2coloutput, _Concat({b, ic * kdv, negone, h, w}, 0));
// Merge od and batch
value = _Transpose(value, {0, 2, 1, 3, 4});
value = _Reshape(value, _Concat({negone, ic, kdv, h, w}, 0));
std::unique_ptr<OpT> op(new OpT);
op->type = OpType_Unpack;
auto axisParam = new AxisT;
axisParam->axis = 2;
op->main.type = OpParameter_Axis;
op->main.value = axisParam;
EXPRP packexpr = Expr::create(std::move(op), {value}, kernelD);
convInputs.resize(kernelD);
for (int i=0; i<kernelD; ++i) {
convInputs[i] = Variable::create(packexpr, i);
}
}
// Make Conv
std::vector<VARP> convOutputs(kernelD);
for (int kd=0; kd<convInputs.size(); ++kd) {
std::unique_ptr<OpT> op(new OpT);
op->type = OpType_Convolution;
op->main.type = OpParameter_Convolution2D;
op->main.value = new Convolution2DT;
auto conv2D = op->main.AsConvolution2D();
conv2D->common.reset(new Convolution2DCommonT);
// Copy common
auto common2d = conv2D->common.get();
common2d->inputCount = common->inputCount;
common2d->outputCount = common->outputCount;
common2d->hasOutputShape = common->hasOutputShape;
common2d->dilateX = common->dilates[2];
common2d->dilateY = common->dilates[1];
common2d->strideX = common->strides[2];
common2d->strideY = common->strides[1];
common2d->pads = {common->pads[1], common->pads[2], common->pads[4], common->pads[5]};
common2d->kernelX = common->kernels[2];
common2d->kernelY = common->kernels[1];
common2d->group = common->group;
common2d->padMode = common->padMode;
// Split Weight
int weightGroupSize = inputChannel*outputChannel / common->group;
conv2D->weight.resize(kernelH * kernelW * weightGroupSize);
for (int i=0; i<weightGroupSize; ++i) {
::memcpy(conv2D->weight.data() + kernelH * kernelW * i, weightPtr + i * kernelD * kernelH * kernelW + kd * kernelH * kernelW, kernelH * kernelW * sizeof(float));
}
conv2D->bias.resize(outputChannel);
::memset(conv2D->bias.data(), 0, outputChannel * sizeof(float));
if (kd == kernelD - 1) {
::memcpy(conv2D->bias.data(), biasDataPtr, outputChannel * sizeof(float));
}
auto convExpr = Expr::create(std::move(op), {convInputs[kd]}, 1);
convOutputs[kd] = Variable::create(convExpr);
convOutputs[kd]->setName(expr->name() + "__" + std::to_string(kd));
}
VARP output;
if (kernelD > 1) {
std::unique_ptr<OpT> op(new OpT);
op->type = OpType_Eltwise;
op->main.type = OpParameter_Eltwise;
op->main.value = new EltwiseT;
op->main.AsEltwise()->type = EltwiseType_SUM;
auto eltExpr = Expr::create(std::move(op), convOutputs);
output = Variable::create(eltExpr);
} else {
output = convOutputs[0];
}
if (common->relu) {
output = _Relu(output);
} else if (common->relu6) {
output = _Relu6(output);
}
// Split od and batch
sx = _Shape(output, true);
w = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(3), {0}), one);
h = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(2), {0}), one);
auto oc = _Unsqueeze(_Scalar<int>(outputChannel), {0});
output = _Reshape(output, _Concat({b, negone, oc, h, w}, 0));
output = _Transpose(output, {0, 2, 1, 3, 4});
output->expr().first->setName(expr->name());
return output->expr().first;
}
static EXPRP _transformConvTranspose3DWithDeconvolution(EXPRP expr) {
std::unique_ptr<MNN::OpT> originOp(expr->get()->UnPack());
auto common = originOp->main.AsConvolution3D()->common.get();
auto weightPtr = originOp->main.AsConvolution3D()->weight.data();
auto biasDataPtr = originOp->main.AsConvolution3D()->bias.data();
auto input = expr->inputs()[0];
// batch, ic, D, H, W -> batch*D, ic, H, W
auto one = _Unsqueeze(_Scalar<int32_t>(1), {0});
auto negone = _Unsqueeze(_Scalar<int>(-1), {0});
auto sx = _Shape(input, true);
auto kernelD = common->kernels[0];
auto kernelH = common->kernels[1];
auto kernelW = common->kernels[2];
auto inputChannel = common->inputCount;
auto outputChannel = common->outputCount;
auto kdv = _Unsqueeze(_Scalar<int>(kernelD), {0});
auto w = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(4), {0}), one);
auto h = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(3), {0}), one);
auto d = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(2), {0}), one);
auto ic = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(1), {0}), one);
auto b = _Slice(sx, _Unsqueeze(_Scalar<int32_t>(0), {0}), one);
input = _Transpose(input, {0, 2, 1, 3, 4});
input = _Reshape(input, _Concat({negone, ic, h, w}, 0));
input->setName(expr->name() + "_gemminput");
// Compute GEMM
std::vector<VARP> convOutputs(kernelD);
for (int kd=0; kd<kernelD; ++kd) {
std::unique_ptr<MNN::OpT> op(new MNN::OpT);
op->type = OpType_Deconvolution;
op->main.type = OpParameter_Convolution2D;
op->main.value = new Convolution2DT;
auto conv2D = op->main.AsConvolution2D();
conv2D->common.reset(new Convolution2DCommonT);
// Copy common
auto common2d = conv2D->common.get();
common2d->inputCount = common->inputCount;
common2d->outputCount = common->outputCount;
common2d->hasOutputShape = common->hasOutputShape;
common2d->dilateX = common->dilates[2];
common2d->dilateY = common->dilates[1];
common2d->strideX = common->strides[2];
common2d->strideY = common->strides[1];
common2d->pads = {common->pads[1], common->pads[2], common->pads[4], common->pads[5]};
if (!common->outPads.empty()) {
common2d->outPads = {common->outPads[1], common->outPads[2], common->outPads[4], common->outPads[5]};
}
common2d->kernelX = common->kernels[2];
common2d->kernelY = common->kernels[1];
common2d->group = common->group;
common2d->padMode = common->padMode;
// Split Weight
int weightGroupSize = inputChannel*outputChannel / common->group;
conv2D->weight.resize(kernelH * kernelW * weightGroupSize);
for (int i=0; i<weightGroupSize; ++i) {
::memcpy(conv2D->weight.data() + kernelH * kernelW * i, weightPtr + i * kernelD * kernelH * kernelW + kd * kernelH * kernelW, kernelH * kernelW * sizeof(float));
}
conv2D->bias.resize(outputChannel);
::memset(conv2D->bias.data(), 0, outputChannel * sizeof(float));
if (kd == kernelD - 1) {
::memcpy(conv2D->bias.data(), biasDataPtr, outputChannel * sizeof(float));
}
auto convExpr = Expr::create(std::move(op), {input}, 1);
convOutputs[kd] = Variable::create(convExpr);
convOutputs[kd]->setName(expr->name() + "__" + std::to_string(kd));
}
auto shapeOutput = _Shape(convOutputs[0], true);
auto ow = _Slice(shapeOutput, _Unsqueeze(_Scalar<int32_t>(3), {0}), one);
auto oh = _Slice(shapeOutput, _Unsqueeze(_Scalar<int32_t>(2), {0}), one);
auto oc = _Unsqueeze(_Scalar<int>(common->outputCount), {0});
auto col2ImInput = _Stack(convOutputs, 0);
// kernelD, batch * D, oc, oh, ow -> batch, oc, kernelD, D, oh, ow
// batch, oc, kernelD, D, oh, ow -> batch, oc * kernelD, D * oh * ow
col2ImInput = _Reshape(col2ImInput, _Concat({kdv, b, negone, oc, oh, ow}, 0));
col2ImInput = _Transpose(col2ImInput, {1, 3, 0, 2, 4, 5});
col2ImInput = _Reshape(col2ImInput, _Concat({b, oc*kdv, negone}, 0));
col2ImInput->setName(expr->name() + "_col2iminput");
// Col2Im
// output_width = (input_width - 1) * sW + dW * (kW - 1) + 1 - layer->pads()->data()[1] - layer->pads()->data()[3] + output_pad;
int outputPad = 0;
if (!common->outPads.empty()) {
outputPad = common->outPads[0];
}
auto outputD = (d - _Scalar<int>(1)) * _Scalar<int>(common->strides[0]) + _Scalar<int>(common->dilates[0] * (common->kernels[0]-1) + 1 - common->pads[0] - common->pads[3] + outputPad);
auto outputHW = _Concat({outputD, oh * ow}, 0);
auto col2ImOutput = _Col2Im(col2ImInput, outputHW, {1, kernelD}, {1, common->dilates[0]}, {common->pads[0], 0, common->pads[3], 0}, {1, common->strides[0]});
col2ImOutput->setName(expr->name() + "_col2imoutput");
auto output = _Reshape(col2ImOutput, _Concat({b, oc, negone, oh, ow}, 0));
if (common->relu) {
output = _Relu(output);
} else if (common->relu6) {
output = _Relu6(output);
}
return output->expr().first;
}
static auto gRegister = []() {
{
auto compare = [](EXPRP expr) {
if (nullptr == expr->get()) {
return false;
}
if (expr->get()->type() != OpType_Convolution3D) {
return false;
}
return expr->get()->type() == OpType_Convolution3D && expr->inputs().size() == 1;
};
auto modify = [](EXPRP expr) {
auto newExpr = _transformConv3DWithConv2D(expr);
newExpr->setName(expr->name());
Expr::replace(expr, newExpr);
return true;
};
TemplateMerge::getInstance("Merge").insertTemplate("Convolution3DTurn2D", compare, modify, PASS_PRIORITY_MIDDLE);
}
{
auto compare = [](EXPRP expr) {
if (nullptr == expr->get()) {
return false;
}
if (expr->get()->type() != OpType_ConvTranspose3D) {
return false;
}
return expr->inputs().size() <= 2;
};
auto modify = [](EXPRP expr) {
auto newExpr = _transformConvTranspose3DWithDeconvolution(expr);
newExpr->setName(expr->name());
Expr::replace(expr, newExpr);
return true;
};
TemplateMerge::getInstance("Merge").insertTemplate("ConvolutionTranspose3DTurn2D", compare, modify, PASS_PRIORITY_MIDDLE);
}
return true;
}();
}
} // namespace MNN