haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add pre-commit action

Browse Source
This commit is contained in:
wang-xinyu 2024-03-19 19:09:02 +08:00
parent f159a101bd
commit f768d67fbc
2 changed files with 147 additions and 93 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
.github/workflow/pre-commit.yml vendored Normal file
View File

@ -0,0 +1,14 @@
name: pre-commit
on:
pull_request:
push:
branches: [main]
jobs:
pre-commit:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
- uses: actions/setup-python@v3
- uses: pre-commit/action@v3.0.1

226
yolov9/src/block.cpp
View File

@ -3,15 +3,15 @@
#include "config.h"
#include "yololayer.h"
#include <iostream>
#include <fstream>
#include <map>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iostream>
#include <map>
#include <numeric>
#include <vector>
#include <algorithm>
using namespace nvinfer1;
@ -19,12 +19,14 @@ using namespace nvinfer1;
// [type] [size] <data x size in hex>
void PrintDim(const ILayer* layer, std::string log) {
Dims dim = layer->getOutput(0)->getDimensions();
std::cout << log <<": "<<"\t\t\t\t";
std::cout << log << ": "
<< "\t\t\t\t";
for (int i = 0; i < dim.nbDims; i++) {
std::cout << dim.d[i] << " ";
}
std::cout << std::endl;
}
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
@ -39,7 +41,7 @@ std::map<std::string, Weights> loadWeights(const std::string file) {
assert(count > 0 && "Invalid weight map file.");
while (count--) {
Weights wt{ DataType::kFLOAT, nullptr, 0 };
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
@ -66,15 +68,17 @@ int get_width(int x, float gw, int divisor) {
}
int get_depth(int x, float gd) {
if (x == 1) return 1;
if (x == 1)
return 1;
int r = round(x * gd);
if (x * gd - int(x * gd) == 0.5 && (int(x * gd) % 2) == 0) {
--r;
}
return std::max<int>(r, 1);
}
static nvinfer1::IScaleLayer* addBatchNorm2d(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
nvinfer1::ITensor& input, std::string lname, float eps) {
static nvinfer1::IScaleLayer* addBatchNorm2d(nvinfer1::INetworkDefinition* network,
std::map<std::string, nvinfer1::Weights> weightMap,
nvinfer1::ITensor& input, std::string lname, float eps) {
float* gamma = (float*)weightMap[lname + ".weight"].values;
float* beta = (float*)weightMap[lname + ".bias"].values;
float* mean = (float*)weightMap[lname + ".running_mean"].values;
@ -82,13 +86,13 @@ nvinfer1::ITensor& input, std::string lname, float eps) {
int len = weightMap[lname + ".running_var"].count;
float* scval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for(int i = 0; i < len; i++) {
for (int i = 0; i < len; i++) {
scval[i] = gamma[i] / sqrt(var[i] + eps);
}
nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, scval, len};
float* shval = reinterpret_cast<float*>(malloc(sizeof(float) * len));
for(int i = 0; i < len; i++) {
for (int i = 0; i < len; i++) {
shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
}
nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, shval, len};
@ -97,7 +101,7 @@ nvinfer1::ITensor& input, std::string lname, float eps) {
for (int i = 0; i < len; i++) {
pval[i] = 1.0;
}
nvinfer1::Weights power{ nvinfer1::DataType::kFLOAT, pval, len };
nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, pval, len};
weightMap[lname + ".scale"] = scale;
weightMap[lname + ".shift"] = shift;
weightMap[lname + ".power"] = power;
@ -105,9 +109,11 @@ nvinfer1::ITensor& input, std::string lname, float eps) {
assert(output);
return output;
}
nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g) {
nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap,
nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g) {
nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0};
nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty);
nvinfer1::IConvolutionLayer* conv =
network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty);
assert(conv);
conv->setStrideNd(nvinfer1::DimsHW{s, s});
conv->setPaddingNd(nvinfer1::DimsHW{p, p});
@ -121,9 +127,12 @@ nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map<std
assert(ew);
return ew;
}
nvinfer1::ILayer* convBnNoAct(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g) {
nvinfer1::ILayer* convBnNoAct(nvinfer1::INetworkDefinition* network,
std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int ch,
int k, int s, int p, std::string lname, int g) {
nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0};
nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty);
nvinfer1::IConvolutionLayer* conv =
network->addConvolutionNd(input, ch, nvinfer1::DimsHW{k, k}, weightMap[lname + ".conv.weight"], bias_empty);
assert(conv);
conv->setStrideNd(nvinfer1::DimsHW{s, s});
conv->setPaddingNd(nvinfer1::DimsHW{p, p});
@ -138,27 +147,31 @@ std::vector<std::vector<float>> getAnchors(std::map<std::string, Weights>& weigh
Weights wts = weightMap[lname + ".anchor_grid"];
int anchor_len = kNumAnchor * 2;
for (int i = 0; i < wts.count / anchor_len; i++) {
auto *p = (const float*)wts.values + i * anchor_len;
auto* p = (const float*)wts.values + i * anchor_len;
std::vector<float> anchor(p, p + anchor_len);
anchors.push_back(anchor);
}
return anchors;
}
ILayer* RepConvN(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int k, int s, int p, int g, int d, bool act, bool bn, bool deploy, std::string lname) {
ILayer* RepConvN(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2,
int k, int s, int p, int g, int d, bool act, bool bn, bool deploy, std::string lname) {
assert(k == 3 && p == 1);
ILayer* conv1 = convBnNoAct(network, weightMap, input, c2, k, s, p, lname + ".conv1", g);
ILayer* conv2 = convBnNoAct(network, weightMap, input, c2, 1, s, p - k / 2, lname + ".conv2", g);
ILayer* ew0 = network->addElementWise(*conv1->getOutput(0), *conv2->getOutput(0), ElementWiseOperation::kSUM);
nvinfer1::IActivationLayer* sigmoid = network->addActivation(*ew0->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
nvinfer1::IActivationLayer* sigmoid =
network->addActivation(*ew0->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
assert(sigmoid);
auto ew = network->addElementWise(*ew0->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD);
auto ew =
network->addElementWise(*ew0->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD);
assert(ew);
return ew;
}
ILayer* RepNBottleneck(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, bool shortcut, int k, int g, float e, std::string lname) {
ILayer* RepNBottleneck(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1,
int c2, bool shortcut, int k, int g, float e, std::string lname) {
int c_ = int(c2 * e);
assert(k == 3 && "RepVGG only support kernel size 3");
auto cv1 = RepConvN(network, weightMap, input, c1, c_, k, 1, 1, g, 1, true, false, false, lname + ".cv1");
@ -170,174 +183,197 @@ ILayer* RepNBottleneck(INetworkDefinition *network, std::map<std::string, Weight
return cv2;
}
ILayer* RepNCSP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input,
int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) {
ILayer* RepNCSP(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2,
int n, bool shortcut, int g, float e, std::string lname) {
int c_ = int(c2 * e);
auto cv1 = convBnSiLU(network, weightMap, input, c_, 1, 1, 0, lname + ".cv1", 1);
ILayer* m = cv1;
for(int i = 0; i < n; i++) {
m = RepNBottleneck(network, weightMap, *m->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m." + std::to_string(i));
for (int i = 0; i < n; i++) {
m = RepNBottleneck(network, weightMap, *m->getOutput(0), c_, c_, shortcut, 3, g, 1.0,
lname + ".m." + std::to_string(i));
}
// auto m_0 = RepNBottleneck(network, weightMap, *cv1->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m.0");
// auto m_0 = RepNBottleneck(network, weightMap, *cv1->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m.0");
// auto m_1 = RepNBottleneck(network, weightMap, *m_0->getOutput(0), c_, c_, shortcut, 3, g, 1.0, lname + ".m.1");
auto cv2 = convBnSiLU(network, weightMap, input, c_, 1, 1, 0, lname + ".cv2", 1);
ITensor* inputTensors[] = { m->getOutput(0), cv2->getOutput(0) };
ITensor* inputTensors[] = {m->getOutput(0), cv2->getOutput(0)};
auto cat = network->addConcatenation(inputTensors, 2);
auto cv3 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv3", 1);
return cv3;
}
ILayer* RepNCSPELAN4(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int c3, int c4, int c5, std::string lname) {
ILayer* RepNCSPELAN4(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1,
int c2, int c3, int c4, int c5, std::string lname) {
auto cv1 = convBnSiLU(network, weightMap, input, c3, 1, 1, 0, lname + ".cv1", 1);
// 将cv1的输出分成两部分 chunk(2, 1)
nvinfer1::Dims d = cv1->getOutput(0)->getDimensions();
nvinfer1::ISliceLayer* split1 = network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{0,0,0}, nvinfer1::Dims3{d.d[0]/2, d.d[1], d.d[2]}, nvinfer1::Dims3{1,1,1});
nvinfer1::ISliceLayer* split2 = network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{d.d[0]/2,0,0}, nvinfer1::Dims3{d.d[0]/2, d.d[1], d.d[2]}, nvinfer1::Dims3{1,1,1});
nvinfer1::ISliceLayer* split1 =
network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{0, 0, 0}, nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]},
nvinfer1::Dims3{1, 1, 1});
nvinfer1::ISliceLayer* split2 =
network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{d.d[0] / 2, 0, 0},
nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, nvinfer1::Dims3{1, 1, 1});
auto cv2_0 = RepNCSP(network, weightMap, *split2->getOutput(0), c3/2, c4, c5, true, 1, 0.5, lname + ".cv2.0");
auto cv2_0 = RepNCSP(network, weightMap, *split2->getOutput(0), c3 / 2, c4, c5, true, 1, 0.5, lname + ".cv2.0");
auto cv2_1 = convBnSiLU(network, weightMap, *cv2_0->getOutput(0), c4, 3, 1, 1, lname + ".cv2.1", 1);
auto cv3_0 = RepNCSP(network, weightMap, *cv2_1->getOutput(0), c4, c4, c5, true, 1, 0.5, lname + ".cv3.0");
auto cv3_1 = convBnSiLU(network, weightMap, *cv3_0->getOutput(0), c4, 3, 1, 1, lname + ".cv3.1", 1);
ITensor* inputTensors[] = { split1->getOutput(0), split2->getOutput(0), cv2_1->getOutput(0), cv3_1->getOutput(0) };
ITensor* inputTensors[] = {split1->getOutput(0), split2->getOutput(0), cv2_1->getOutput(0), cv3_1->getOutput(0)};
auto cat = network->addConcatenation(inputTensors, 4);
auto cv4 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv4", 1);
return cv4;
}
ILayer* ADown(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2, std::string lname) {
ILayer* ADown(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2,
std::string lname) {
int c_ = c2 / 2;
auto pool = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{ 2, 2 });
pool->setStrideNd(DimsHW{ 1, 1 });
pool->setPaddingNd(DimsHW{ 0, 0 });
auto pool = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{2, 2});
pool->setStrideNd(DimsHW{1, 1});
pool->setPaddingNd(DimsHW{0, 0});
nvinfer1::Dims d = pool->getOutput(0)->getDimensions();
nvinfer1::ISliceLayer* split1 = network->addSlice(*pool->getOutput(0), nvinfer1::Dims3{ 0, 0, 0}, nvinfer1::Dims3{ d.d[0]/2, d.d[1], d.d[2] }, nvinfer1::Dims3{ 1, 1, 1 });
nvinfer1::ISliceLayer* split2 = network->addSlice(*pool->getOutput(0), nvinfer1::Dims3{ d.d[0]/2, 0, 0}, nvinfer1::Dims3{ d.d[0]/2, d.d[1], d.d[2] }, nvinfer1::Dims3{ 1, 1, 1 });
nvinfer1::ISliceLayer* split1 =
network->addSlice(*pool->getOutput(0), nvinfer1::Dims3{0, 0, 0},
nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, nvinfer1::Dims3{1, 1, 1});
nvinfer1::ISliceLayer* split2 =
network->addSlice(*pool->getOutput(0), nvinfer1::Dims3{d.d[0] / 2, 0, 0},
nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, nvinfer1::Dims3{1, 1, 1});
// auto chunklayer = layer_split(1, pool->getOutput(0), network);
auto cv1 = convBnSiLU(network, weightMap, *split1->getOutput(0), c_, 3, 2, 1, lname + ".cv1", 1);
auto pool2 = network->addPoolingNd(*split2->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 });
pool2->setStrideNd(DimsHW{ 2, 2 });
pool2->setPaddingNd(DimsHW{ 1, 1 });
auto pool2 = network->addPoolingNd(*split2->getOutput(0), PoolingType::kMAX, DimsHW{3, 3});
pool2->setStrideNd(DimsHW{2, 2});
pool2->setPaddingNd(DimsHW{1, 1});
auto cv2 = convBnSiLU(network, weightMap, *pool2->getOutput(0), c_, 1, 1, 0, lname + ".cv2", 1);
ITensor* inputTensors[] = { cv1->getOutput(0), cv2->getOutput(0) };
ITensor* inputTensors[] = {cv1->getOutput(0), cv2->getOutput(0)};
auto cat = network->addConcatenation(inputTensors, 2);
return cat;
}
std::vector<ILayer*> CBLinear(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, std::vector<int> c2s, int k, int s, int p, int g, std::string lname) {
std::vector<ILayer*> CBLinear(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input,
std::vector<int> c2s, int k, int s, int p, int g, std::string lname) {
IConvolutionLayer* conv1 = network->addConvolutionNd(input, std::accumulate(c2s.begin(), c2s.end(), 0), DimsHW{ k, k }, weightMap[lname + ".conv.weight"], weightMap[lname + ".conv.bias"]);
IConvolutionLayer* conv1 =
network->addConvolutionNd(input, std::accumulate(c2s.begin(), c2s.end(), 0), DimsHW{k, k},
weightMap[lname + ".conv.weight"], weightMap[lname + ".conv.bias"]);
assert(conv1);
conv1->setName((lname + ".conv").c_str());
conv1->setStrideNd(DimsHW{ s, s });
conv1->setPaddingNd(DimsHW{ p, p });
conv1->setStrideNd(DimsHW{s, s});
conv1->setPaddingNd(DimsHW{p, p});
int h = input.getDimensions().d[1];
int w = input.getDimensions().d[2];
std::vector<ILayer*> slices(c2s.size());
int start = 0;
for(int i = 0; i < c2s.size(); i++) {
slices[i] = network->addSlice(*conv1->getOutput(0), Dims3{ start, 0, 0 }, Dims3{ c2s[i], h, w }, Dims3{ 1, 1, 1 });
for (int i = 0; i < c2s.size(); i++) {
slices[i] = network->addSlice(*conv1->getOutput(0), Dims3{start, 0, 0}, Dims3{c2s[i], h, w}, Dims3{1, 1, 1});
start += c2s[i];
}
return slices;
}
ILayer* CBFuse(INetworkDefinition *network, std::vector<std::vector<ILayer*>> input, std::vector<int> idx, std::vector<int> strides) {
ILayer* CBFuse(INetworkDefinition* network, std::vector<std::vector<ILayer*>> input, std::vector<int> idx,
std::vector<int> strides) {
ILayer** res = new ILayer*[input.size()];
res[input.size()-1] = input[input.size()-1][0];
res[input.size() - 1] = input[input.size() - 1][0];
for(int i = input.size()-2; i >= 0; i--) {
for (int i = input.size() - 2; i >= 0; i--) {
auto upsample = network->addResize(*input[i][idx[i]]->getOutput(0));
upsample->setResizeMode(ResizeMode::kNEAREST);
const float scales[] = { 1, strides[i] / strides[strides.size() - 1], strides[i] / strides[strides.size() - 1] };
const float scales[] = {1, strides[i] / strides[strides.size() - 1], strides[i] / strides[strides.size() - 1]};
upsample->setScales(scales, 3);
res[i] = upsample;
}
for(int i = 1; i < input.size(); i++) {
for (int i = 1; i < input.size(); i++) {
auto ew = network->addElementWise(*res[0]->getOutput(0), *res[i]->getOutput(0), ElementWiseOperation::kSUM);
res[0] = ew;
}
return res[0];
}
ILayer* SP(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int k, int s) {
ILayer* SP(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int k, int s) {
int p = k / 2;
auto pool = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{ k, k });
pool->setPaddingNd(DimsHW{ p, p });
pool->setStrideNd(DimsHW{ s, s });
auto pool = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{k, k});
pool->setPaddingNd(DimsHW{p, p});
pool->setStrideNd(DimsHW{s, s});
return pool;
}
ILayer* SPPELAN(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2, int c3, std::string lname) {
ILayer* SPPELAN(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c1, int c2,
int c3, std::string lname) {
auto cv1 = convBnSiLU(network, weightMap, input, c3, 1, 1, 0, lname + ".cv1", 1);
auto cv2 = SP(network, weightMap, *cv1->getOutput(0), 5, 1);
auto cv3 = SP(network, weightMap, *cv2->getOutput(0), 5, 1);
auto cv4 = SP(network, weightMap, *cv3->getOutput(0), 5, 1);
ITensor* inputTensors[] = { cv1->getOutput(0), cv2->getOutput(0), cv3->getOutput(0), cv4->getOutput(0) };
ITensor* inputTensors[] = {cv1->getOutput(0), cv2->getOutput(0), cv3->getOutput(0), cv4->getOutput(0)};
auto cat = network->addConcatenation(inputTensors, 4);
auto cv5 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv5", 1);
return cv5;
}
ILayer* DetectBbox_Conv(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2, int reg_max, std::string lname) {
ILayer* DetectBbox_Conv(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2,
int reg_max, std::string lname) {
auto cv_0 = convBnSiLU(network, weightMap, input, c2, 3, 1, 1, lname + ".0", 1);
auto cv_1 = convBnSiLU(network, weightMap, *cv_0->getOutput(0), c2, 3, 1, 1, lname + ".1", 4);
auto cv_2 = network->addConvolutionNd(*cv_1->getOutput(0), reg_max*4, DimsHW{ 1, 1 }, weightMap[lname + ".2.weight"], weightMap[lname + ".2.bias"]);
auto cv_2 = network->addConvolutionNd(*cv_1->getOutput(0), reg_max * 4, DimsHW{1, 1},
weightMap[lname + ".2.weight"], weightMap[lname + ".2.bias"]);
cv_2->setName((lname + ".conv").c_str());
cv_2->setStrideNd(DimsHW{ 1, 1 });
cv_2->setPaddingNd(DimsHW{ 0, 0 });
cv_2->setStrideNd(DimsHW{1, 1});
cv_2->setPaddingNd(DimsHW{0, 0});
cv_2->setNbGroups(4);
return cv_2;
}
ILayer* DetectCls_Conv(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2, int cls, std::string lname) {
ILayer* DetectCls_Conv(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int c2,
int cls, std::string lname) {
auto cv_0 = convBnSiLU(network, weightMap, input, c2, 3, 1, 1, lname + ".0", 1);
auto cv_1 = convBnSiLU(network, weightMap, *cv_0->getOutput(0), c2, 3, 1, 1, lname + ".1", 1);
auto cv_2 = network->addConvolutionNd(*cv_1->getOutput(0), cls, DimsHW{ 1, 1 }, weightMap[lname + ".2.weight"], weightMap[lname + ".2.bias"]);
auto cv_2 = network->addConvolutionNd(*cv_1->getOutput(0), cls, DimsHW{1, 1}, weightMap[lname + ".2.weight"],
weightMap[lname + ".2.bias"]);
cv_2->setName((lname + ".conv").c_str());
cv_2->setStrideNd(DimsHW{ 1, 1 });
cv_2->setPaddingNd(DimsHW{ 0, 0 });
cv_2->setStrideNd(DimsHW{1, 1});
cv_2->setPaddingNd(DimsHW{0, 0});
return cv_2;
}
nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname){
nvinfer1::IShuffleLayer* DFL(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname) {
auto dim = input.getDimensions();
int c = dim.d[0];
int grid = dim.d[1] * dim.d[2];
int split_num = c / ch;
nvinfer1::IShuffleLayer* shuffle1 = network->addShuffle(input);
shuffle1->setReshapeDimensions(nvinfer1::Dims3{ split_num, ch, grid });
shuffle1->setSecondTranspose(nvinfer1::Permutation{ 1, 0, 2 });
shuffle1->setReshapeDimensions(nvinfer1::Dims3{split_num, ch, grid});
shuffle1->setSecondTranspose(nvinfer1::Permutation{1, 0, 2});
nvinfer1::ISoftMaxLayer* softmax = network->addSoftMax(*shuffle1->getOutput(0));
nvinfer1::Weights bias_empty{ nvinfer1::DataType::kFLOAT, nullptr, 0 };
nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(*softmax->getOutput(0), 1, nvinfer1::DimsHW{1, 1}, weightMap[lname + ".conv.weight"], bias_empty);
conv->setStrideNd(nvinfer1::DimsHW{ s, s });
conv->setPaddingNd(nvinfer1::DimsHW{ p, p });
nvinfer1::Weights bias_empty{nvinfer1::DataType::kFLOAT, nullptr, 0};
nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(*softmax->getOutput(0), 1, nvinfer1::DimsHW{1, 1},
weightMap[lname + ".conv.weight"], bias_empty);
conv->setStrideNd(nvinfer1::DimsHW{s, s});
conv->setPaddingNd(nvinfer1::DimsHW{p, p});
nvinfer1::IShuffleLayer* shuffle2 = network->addShuffle(*conv->getOutput(0));
shuffle2->setReshapeDimensions(nvinfer1::Dims2{ 4, grid });
shuffle2->setReshapeDimensions(nvinfer1::Dims2{4, grid});
return shuffle2;
}
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation) {
nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition* network,
std::vector<nvinfer1::IConcatenationLayer*> dets, bool is_segmentation) {
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
nvinfer1::PluginField plugin_fields[1];
int netinfo[5] = { kNumClass, kInputW, kInputH, kMaxNumOutputBbox, is_segmentation };
int netinfo[5] = {kNumClass, kInputW, kInputH, kMaxNumOutputBbox, is_segmentation};
plugin_fields[0].data = netinfo;
plugin_fields[0].length = 5;
plugin_fields[0].name = "netinfo";
@ -346,16 +382,18 @@ nvinfer1::IPluginV2Layer* addYoLoLayer(nvinfer1::INetworkDefinition *network, st
nvinfer1::PluginFieldCollection plugin_data;
plugin_data.nbFields = 1;
plugin_data.fields = plugin_fields;
nvinfer1::IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data);
nvinfer1::IPluginV2* plugin_obj = creator->createPlugin("yololayer", &plugin_data);
std::vector<nvinfer1::ITensor*> input_tensors;
for (auto det: dets) {
for (auto det : dets) {
input_tensors.push_back(det->getOutput(0));
}
auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(), *plugin_obj);
return yolo;
}
std::vector<IConcatenationLayer*> DualDDetect(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::vector<ILayer*> dets, int cls, std::vector<int> ch, std::string lname) {
std::vector<IConcatenationLayer*> DualDDetect(INetworkDefinition* network, std::map<std::string, Weights>& weightMap,
std::vector<ILayer*> dets, int cls, std::vector<int> ch,
std::string lname) {
int c2 = std::max(int(ch[0] / 4), int(16 * 4));
int c3 = std::max(ch[0], std::min(cls * 2, 128));
int reg_max = 16;
@ -365,12 +403,14 @@ std::vector<IConcatenationLayer*> DualDDetect(INetworkDefinition *network, std::
for (int i = 0; i < dets.size(); i++) {
// Conv(x, c2, 3), Conv(c2, c2, 3, g=4), nn.Conv2d(c2, 4 * self.reg_max, 1, groups=4)
bboxlayers.push_back(DetectBbox_Conv(network, weightMap, *dets[i]->getOutput(0), c2, reg_max, lname + ".cv2." + std::to_string(i)));
bboxlayers.push_back(DetectBbox_Conv(network, weightMap, *dets[i]->getOutput(0), c2, reg_max,
lname + ".cv2." + std::to_string(i)));
// Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, self.nc, 1)
auto cls_layer = DetectCls_Conv(network, weightMap, *dets[i]->getOutput(0), c3, cls, lname + ".cv3." + std::to_string(i));
auto cls_layer = DetectCls_Conv(network, weightMap, *dets[i]->getOutput(0), c3, cls,
lname + ".cv3." + std::to_string(i));
auto dim = cls_layer->getOutput(0)->getDimensions();
nvinfer1::IShuffleLayer* shuffle = network->addShuffle(*cls_layer->getOutput(0));
shuffle->setReshapeDimensions(nvinfer1::Dims2{ kNumClass, dim.d[1] * dim.d[2] });
shuffle->setReshapeDimensions(nvinfer1::Dims2{kNumClass, dim.d[1] * dim.d[2]});
clslayers.push_back(shuffle);
}
@ -378,8 +418,8 @@ std::vector<IConcatenationLayer*> DualDDetect(INetworkDefinition *network, std::
for (int i = 0; i < dets.size(); i++) {
// softmax 16*4, w, h => 16, 4, w, h
auto loc = DFL(network, weightMap, *bboxlayers[i]->getOutput(0), 16, 1, 1, 0, lname + ".dfl");
nvinfer1::ITensor* inputTensor[] = { loc->getOutput(0), clslayers[i]->getOutput(0) };
ret.push_back(network->addConcatenation(inputTensor, 2));
nvinfer1::ITensor* inputTensor[] = {loc->getOutput(0), clslayers[i]->getOutput(0)};
ret.push_back(network->addConcatenation(inputTensor, 2));
}
return ret;
}