327 lines
14 KiB
C++
327 lines
14 KiB
C++
#include "common.hpp"
|
|
#include "logging.h"
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <map>
|
|
#include <sstream>
|
|
#include <vector>
|
|
#include <chrono>
|
|
|
|
#define USE_FP32
|
|
|
|
static Logger gLogger;
|
|
|
|
const char *INPUT_BLOB_NAME = "data";
|
|
const char *OUTPUT_BLOB_NAME = "output";
|
|
static const int bs = 1;
|
|
static const int channels = 96;
|
|
static const int ch = 3;
|
|
static const int INPUT_H = 576;
|
|
static const int INPUT_W = 576;
|
|
static const int NUM_CLASSES = 15;
|
|
static const int outputSize = 576 * 576;
|
|
cudaStream_t m_cudaStream;
|
|
vector<void *> m_bindings;
|
|
IExecutionContext *m_context;
|
|
|
|
ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, DataType dt,std::string wtsPath)
|
|
{
|
|
INetworkDefinition *network = builder->createNetworkV2(0U);
|
|
ITensor *data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ch, INPUT_H, INPUT_W});
|
|
assert(data);
|
|
std::map<std::string, Weights> weightMap = loadWeights(wtsPath);
|
|
ITensor* conv1 = conv(network, weightMap, data, "backbone.patch_embed.proj", channels);
|
|
ITensor* shuffle1 = shuffle_reshapeApermute(network, conv1, Dims2{channels, -1}, Permutation{1, 0}, true);
|
|
ITensor *ln = m_layerNorm(network, weightMap, shuffle1, "backbone.patch_embed.norm");
|
|
debug_print(ln, "ln");
|
|
//layer0
|
|
|
|
ITensor *mask0 = trt_transform_imgMask(network, 147, 7, 3);
|
|
ITensor *blk00 = blk(network, weightMap, ln, mask0, "backbone.layers.0.blocks.0", INPUT_H / 4, channels, 3, 7, 0);
|
|
debug_print(blk00, "blk00");
|
|
ITensor *blk01 = blk(network, weightMap, blk00, mask0, "backbone.layers.0.blocks.1", INPUT_H / 4, channels, 3, 7, 3);
|
|
debug_print(blk01, "blk01");
|
|
ITensor* out0 = m_layerNorm(network, weightMap, blk01, "backbone.norm0");
|
|
out0 = shuffle_reshapeApermute(network, out0, Dims3{INPUT_H / 4, INPUT_H / 4, channels}, Permutation{2, 0, 1}, true);
|
|
ITensor *down_layer0 = downsample(network, weightMap, blk01, "backbone.layers.0.downsample", INPUT_H / 4);
|
|
debug_print(down_layer0, "down_blk1");
|
|
//layer1
|
|
ITensor *mask1 = trt_transform_imgMask(network, 77, 7, 3);
|
|
ITensor *blk10 = blk(network, weightMap, down_layer0, mask1, "backbone.layers.1.blocks.0", INPUT_H / 8, channels * 2, 6, 7, 0);
|
|
debug_print(blk10, "blk10");
|
|
ITensor *blk11 = blk(network, weightMap, blk10, mask1, "backbone.layers.1.blocks.1", INPUT_H / 8, channels * 2, 6, 7, 3);
|
|
debug_print(blk11, "blk11");
|
|
ITensor* out1 = m_layerNorm(network, weightMap, blk11, "backbone.norm1");
|
|
out1 = shuffle_reshapeApermute(network, out1, Dims3{INPUT_H / 8, INPUT_H / 8, channels * 2}, Permutation{2, 0, 1}, true);
|
|
ITensor *down_layer1 = downsample(network, weightMap, blk11, "backbone.layers.1.downsample", INPUT_H / 8);
|
|
debug_print(down_layer1, "down_layer1");
|
|
//layer2
|
|
ITensor *mask2 = trt_transform_imgMask(network, 42, 7, 3);
|
|
ITensor *blk20 = blk(network, weightMap, down_layer1, mask2, "backbone.layers.2.blocks.0", INPUT_H / 16, channels * 4, 12, 7, 0);
|
|
debug_print(blk20, "blk20");
|
|
ITensor *blk21 = blk(network, weightMap, blk20, mask2, "backbone.layers.2.blocks.1", INPUT_H / 16, channels * 4, 12, 7, 3);
|
|
debug_print(blk21, "blk21");
|
|
ITensor *blk22 = blk(network, weightMap, blk21, mask2, "backbone.layers.2.blocks.2", INPUT_H / 16,channels * 4, 12, 7, 0);
|
|
debug_print(blk22, "blk22");
|
|
ITensor *blk23 = blk(network, weightMap, blk22, mask2, "backbone.layers.2.blocks.3", INPUT_H / 16, channels * 4, 12, 7, 3);
|
|
debug_print(blk23, "blk23");
|
|
ITensor *blk24 = blk(network, weightMap, blk23, mask2, "backbone.layers.2.blocks.4", INPUT_H / 16, channels * 4, 12, 7, 0);
|
|
debug_print(blk24, "blk24");
|
|
ITensor *blk25 = blk(network, weightMap, blk24, mask2, "backbone.layers.2.blocks.5", INPUT_H / 16, channels * 4, 12, 7, 3);
|
|
debug_print(blk25, "blk25");
|
|
ITensor* out2 = m_layerNorm(network, weightMap, blk25, "backbone.norm2");
|
|
out2 = shuffle_reshapeApermute(network, out2, Dims3{INPUT_H / 16, INPUT_H / 16, channels * 4}, Permutation{2, 0, 1}, true);
|
|
ITensor *down_layer2 = downsample(network, weightMap, blk25, "backbone.layers.2.downsample", INPUT_H / 16);
|
|
debug_print(down_layer2, "down_layer2");
|
|
//layer3
|
|
ITensor *mask3 = trt_transform_imgMask(network, 21, 7, 3);
|
|
ITensor *blk30 = blk(network, weightMap, down_layer2, mask3, "backbone.layers.3.blocks.0", INPUT_H / 32, channels * 8, 24, 7, 0);
|
|
debug_print(blk30, "blk30");
|
|
ITensor *blk31 = blk(network, weightMap, blk30, mask3, "backbone.layers.3.blocks.1", INPUT_H / 32, channels * 8, 24, 7, 3);
|
|
debug_print(blk31, "blk31");
|
|
ITensor* out3 = m_layerNorm(network, weightMap, blk31, "backbone.norm3");
|
|
out3 = shuffle_reshapeApermute(network, out3, Dims3{INPUT_H / 32, INPUT_H / 32, channels * 8}, Permutation{2, 0, 1}, true);
|
|
ITensor* out[4] = {out0, out1, out2, out3};
|
|
out0 = transform_lateral_conv(network, weightMap, out0, "decode_head.lateral_convs.0"); // 512,INPUT_H/4,INPUT_H/4
|
|
out1 = transform_lateral_conv(network, weightMap, out1, "decode_head.lateral_convs.1"); // 512,INPUT_H/8,INPUT_H/8
|
|
out2 = transform_lateral_conv(network, weightMap, out2, "decode_head.lateral_convs.2"); // 512,INPUT_H/16,INPUT_H/16
|
|
auto psp_out_0 = transform_psp(network, weightMap, out3, "decode_head.psp_modules.0.1", 1);
|
|
auto psp_out_1 = transform_psp(network, weightMap, out3, "decode_head.psp_modules.1.1", 2);
|
|
auto psp_out_2 = transform_psp(network, weightMap, out3, "decode_head.psp_modules.2.1", 3);
|
|
auto psp_out_3 = transform_psp(network, weightMap, out3, "decode_head.psp_modules.3.1", 6);
|
|
ITensor* psp_outs[5] = {out3, psp_out_0, psp_out_1, psp_out_2, psp_out_3};
|
|
auto PSP_outs = network->addConcatenation(psp_outs, 5);
|
|
PSP_outs->setAxis(0);
|
|
debug_print(PSP_outs->getOutput(0), "PSP_outs");
|
|
out3 = transform_lateral_conv(network, weightMap, PSP_outs->getOutput(0), "decode_head.bottleneck", 3, 1, 512); // 512,INPUT_H/32,INPUT_H/32
|
|
debug_print(out3, "out3");
|
|
auto laterals2 = up_Add(network, out3, out2);
|
|
auto laterals1 = up_Add(network, laterals2, out1);
|
|
auto laterals0 = up_Add(network, laterals1, out0);
|
|
auto fpn0 = transform_lateral_conv(network, weightMap, laterals0, "decode_head.fpn_convs.0", 3, 1, 512);
|
|
auto fpn1 = transform_lateral_conv(network, weightMap, laterals1, "decode_head.fpn_convs.1", 3, 1, 512);
|
|
auto fpn2 = transform_lateral_conv(network, weightMap, laterals2, "decode_head.fpn_convs.2", 3, 1, 512);
|
|
fpn1 = resize(network, fpn1,fpn0->getDimensions().d[1]);
|
|
fpn2 = resize(network, fpn2,fpn0->getDimensions().d[1]);
|
|
auto fpn3 = resize(network, out3, fpn0->getDimensions().d[1]);
|
|
ITensor* fpn_outs[4] = {fpn0, fpn1, fpn2, fpn3};
|
|
auto FPN_outs = network->addConcatenation(fpn_outs, 4);
|
|
FPN_outs->setAxis(0);
|
|
debug_print(FPN_outs->getOutput(0), "FPN_outs");
|
|
auto fpn_output = transform_lateral_conv(network, weightMap, FPN_outs->getOutput(0), "decode_head.fpn_bottleneck", 3, 1, 512);
|
|
debug_print(fpn_output, "fpn_output");
|
|
auto seg = network->addConvolutionNd(*fpn_output, NUM_CLASSES, Dims2{1, 1}, weightMap["decode_head.conv_seg.weight"], weightMap["decode_head.conv_seg.bias"]);
|
|
seg->setStrideNd(Dims2{1, 1});
|
|
debug_print(seg->getOutput(0), "seg");
|
|
auto seg_resize = resize(network, seg->getOutput(0), INPUT_H);
|
|
debug_print(seg_resize, "seg_resize");
|
|
auto output = network->addTopK(*seg_resize, TopKOperation::kMAX, 1, 0X01)->getOutput(1);
|
|
debug_print(output, "output");
|
|
|
|
std::cout << "set name out" << std::endl;
|
|
output->setName(OUTPUT_BLOB_NAME);
|
|
network->markOutput(*output);
|
|
builder->setMaxBatchSize(12);
|
|
config->setMaxWorkspaceSize((1 << 30)); // 1G
|
|
#ifdef USE_FP16
|
|
std::cout<< "use fp16"<<std::endl;
|
|
config->setFlag(BuilderFlag::kFP16);
|
|
#endif
|
|
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
|
|
std::cout << "build success!" << std::endl;
|
|
network->destroy();
|
|
|
|
return engine;
|
|
}
|
|
|
|
void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream,std::string wtsPath)
|
|
{
|
|
IBuilder *builder = createInferBuilder(gLogger);
|
|
IBuilderConfig *config = builder->createBuilderConfig();
|
|
ICudaEngine *engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT, wtsPath);
|
|
assert(engine != nullptr);
|
|
(*modelStream) = engine->serialize();
|
|
engine->destroy();
|
|
builder->destroy();
|
|
}
|
|
|
|
void createEng(std::string wtsPath, std::string engine_name)
|
|
{
|
|
char *trtModelStream{nullptr};
|
|
size_t size{0};
|
|
|
|
IHostMemory *modelStream{nullptr};
|
|
APIToModel(bs, &modelStream, wtsPath);
|
|
assert(modelStream != nullptr);
|
|
std::ofstream p(engine_name, std::ios::binary);
|
|
if (!p)
|
|
{
|
|
std::cerr << "could not open plan output file" << std::endl;
|
|
return;
|
|
}
|
|
p.write(reinterpret_cast<const char *>(modelStream->data()), modelStream->size());
|
|
modelStream->destroy();
|
|
std::ifstream file(engine_name, std::ios::binary);
|
|
if (file.good())
|
|
{
|
|
file.seekg(0, file.end);
|
|
size = file.tellg();
|
|
file.seekg(0, file.beg);
|
|
trtModelStream = new char[size];
|
|
assert(trtModelStream);
|
|
file.read(trtModelStream, size);
|
|
file.close();
|
|
}
|
|
}
|
|
|
|
void inference_init(string ENGPath,ICudaEngine *m_engine)
|
|
{
|
|
ifstream cache(ENGPath, ios::binary);
|
|
cache.seekg(0, ios::end);
|
|
const int engSize = cache.tellg();
|
|
cache.seekg(0, ios::beg);
|
|
void *modelMem = malloc(engSize);
|
|
cache.read((char*)modelMem, engSize);
|
|
cache.close();
|
|
IRuntime *runtime = nvinfer1::createInferRuntime(gLogger);
|
|
m_engine = runtime->deserializeCudaEngine(modelMem, engSize);
|
|
runtime->destroy();
|
|
free(modelMem);
|
|
if (!m_engine) {
|
|
cout << "deserialize eng error!" << endl;
|
|
return;
|
|
}
|
|
m_context = m_engine->createExecutionContext();
|
|
if (cudaStreamCreate(&m_cudaStream) != 0) return;
|
|
int bindings = m_engine->getNbBindings();
|
|
if (bindings < 2)
|
|
{
|
|
cout << "Error! the network have one input and one output at least!" << endl;
|
|
return;
|
|
}
|
|
cout << "1111111111111" << endl;
|
|
m_bindings.resize(bindings, nullptr);
|
|
CHECK(cudaMalloc(&m_bindings.at(0), bs * ch * INPUT_H * INPUT_W * sizeof(float)));
|
|
CHECK(cudaMalloc(&m_bindings.at(1), bs * outputSize * 4));
|
|
}
|
|
|
|
void doInference(const float *input, int *output)
|
|
{
|
|
cout << "do infer:" << endl;
|
|
CHECK(cudaMemcpyAsync(m_bindings.at(0), input, bs * ch * INPUT_H * INPUT_W * sizeof(float),
|
|
cudaMemcpyHostToDevice, m_cudaStream));
|
|
|
|
m_context->enqueue(bs, m_bindings.data(), m_cudaStream, nullptr);
|
|
|
|
CHECK(cudaMemcpyAsync(output, m_bindings.at(1), bs * outputSize * 4,
|
|
cudaMemcpyDeviceToHost, m_cudaStream));
|
|
|
|
cudaStreamSynchronize(m_cudaStream);
|
|
}
|
|
|
|
|
|
int main(int argc, char** argv)
|
|
{
|
|
cout << "begin" << endl;
|
|
//string wts = "G:/shaj/trainsform/ktn5n6_29.511.21.8.wts";
|
|
//string eng = "G:/shaj/trainsform/trainsform.eng";
|
|
std::string argv1 = argv[1];
|
|
if (argv1 == "-s") {
|
|
string wts = argv[2];
|
|
string eng = argv[3];
|
|
createEng(wts,eng);
|
|
} else {
|
|
string eng = argv[2];
|
|
|
|
ICudaEngine *m_engine;
|
|
|
|
inference_init(eng,m_engine);
|
|
|
|
vector<cv::Mat> testVal;
|
|
map<string,cv::Mat> dataProb;
|
|
vector<string> imgs;
|
|
cv::Mat img;
|
|
//string pattern_dir = "G:/shaj/trainsform";
|
|
string pattern_dir = argv[3];
|
|
string pattern = pattern_dir+ "/*.bmp";
|
|
vector<cv::String> images_names;
|
|
cv::glob(pattern, images_names, false);
|
|
int i = 0;
|
|
cv::Scalar Mean = cv::Scalar(123.675, 116.28, 103.53);
|
|
cv::Scalar Std = cv::Scalar(58.395, 57.12, 57.375);
|
|
cv::Size size = {INPUT_H,INPUT_W};
|
|
|
|
for (auto image_name: images_names)
|
|
{
|
|
if (i < bs)
|
|
{
|
|
cv::Mat Img = cv::imread(image_name, 1);
|
|
|
|
testVal.push_back(Img);
|
|
cout << image_name << endl;
|
|
imgs.push_back(image_name);
|
|
}
|
|
}
|
|
float *data = new float[bs * ch * INPUT_H * INPUT_W];
|
|
int *output = new int[bs * outputSize];
|
|
|
|
cv::Mat Transed_t = BlobFromImages(testVal, cv::Size{INPUT_H, INPUT_W}, Mean, Std, true, false);
|
|
|
|
memcpy(data, Transed_t.data, bs * ch * INPUT_H * INPUT_W * sizeof(float));
|
|
|
|
//for(int i = 0 ; i< 20; i++){
|
|
|
|
auto start_time = std::chrono::system_clock::now();
|
|
doInference(data, output);
|
|
|
|
auto end_time = std::chrono::system_clock::now();
|
|
float duration;
|
|
duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
|
|
cout << "time:" << duration << endl;
|
|
//}
|
|
// for(int i = 0; i < 100; i++)
|
|
// cout<<i<<":"<<output[i]<<endl;
|
|
|
|
int n = 0;
|
|
int *out = new int[outputSize];
|
|
string outPath = pattern_dir + "/output";
|
|
for (int i = 0; i < testVal.size(); i++)
|
|
{
|
|
cv::Mat img = cv::imread(imgs[i], 1);
|
|
cv::Mat dst;
|
|
cv::resize(img,dst,cv::Size{INPUT_H, INPUT_W});
|
|
//string outPath_n = outPath + "/"+to_string(n) + ".jpg";
|
|
n += 1;
|
|
out = output + i * outputSize;
|
|
for (int i = 0; i < outputSize; i++)
|
|
{
|
|
if (out[i] != 0)
|
|
{
|
|
int w = i % (INPUT_H);
|
|
int h = i / (INPUT_W);
|
|
dst.at<cv::Vec3b>(h, w)[0] = out[i] * 10;
|
|
dst.at<cv::Vec3b>(h, w)[1] = out[i] * 30;
|
|
dst.at<cv::Vec3b>(h, w)[2] = out[i] * 40;
|
|
}
|
|
}
|
|
//cout<<outPath_n<<endl;
|
|
string outPath_result = imgs[i].replace(0, pattern_dir.size(), outPath);
|
|
cout << outPath_result << endl;
|
|
cv::imwrite(outPath_result, dst);
|
|
}
|
|
testVal.clear();
|
|
imgs.clear();
|
|
}
|
|
|
|
m_context->destroy();
|
|
m_engine->destroy();
|
|
for (auto bindings: m_bindings) {
|
|
cudaFree(bindings);
|
|
}
|
|
cudaFree(m_cudaStream);
|
|
|
|
cout << "swin_transform" << endl;
|
|
return 0;
|
|
}
|
|
|