diff --git a/yolov7/README.md b/yolov7/README.md index 3b83933..fdd4876 100644 --- a/yolov7/README.md +++ b/yolov7/README.md @@ -41,14 +41,14 @@ cd {tensorrtx}/yolov7/ // update CLASS_NUM in yololayer.h if your model is trained on custom dataset mkdir build cd build -cp {WongKinYiu}/yolov7/yolov7s.wts {tensorrtx}/yolov7/build +cp {WongKinYiu}/yolov7/yolov7.wts {tensorrtx}/yolov7/build cmake .. make sudo ./yolov7 -s [.wts] [.engine] [t/v7/x/w6/e6/d6/e6e gd gw] // serialize model to plan file sudo ./yolov7 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. -// For example yolov7s +// For example yolov7 sudo ./yolov7 -s yolov7.wts yolov7.engine v7 -sudo ./yolov7 -d yolov7s.engine ../samples +sudo ./yolov7 -d yolov7.engine ../samples ``` 3. check the images generated, as follows. _zidane.jpg and _bus.jpg @@ -82,4 +82,3 @@ python yolov7_trt.py ## More Information See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx) - diff --git a/yolov7/gen_wts.py b/yolov7/gen_wts.py index 29cb65e..9c073a4 100644 --- a/yolov7/gen_wts.py +++ b/yolov7/gen_wts.py @@ -1,16 +1,48 @@ +import sys +import argparse +import os import struct import torch +from utils.torch_utils import select_device -state_dict = torch.load("yolov7.pt", map_location="cpu") -state_dict = state_dict["model"].state_dict() -with open("yolov7.wts", 'w') as f: - f.write("{}\n".format(len(state_dict.keys()))) +def parse_args(): + parser = argparse.ArgumentParser(description='Convert .pt file to .wts') + parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)') + parser.add_argument('-o', '--output', help='Output (.wts) file path (optional)') + args = parser.parse_args() + if not os.path.isfile(args.weights): + raise SystemExit('Invalid input file') + if not args.output: + args.output = os.path.splitext(args.weights)[0] + '.wts' + elif os.path.isdir(args.output): + args.output = os.path.join( + args.output, + os.path.splitext(os.path.basename(args.weights))[0] + '.wts') + return args.weights, args.output - for k, v in state_dict.items(): + +pt_file, wts_file = parse_args() + +# Initialize +device = select_device('cpu') +# Load model +model = torch.load(pt_file, map_location=device)['model'].float() # load to FP32 + +# update anchor_grid info +anchor_grid = model.model[-1].anchors * model.model[-1].stride[...,None,None] +# model.model[-1].anchor_grid = anchor_grid +delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer +model.model[-1].register_buffer("anchor_grid",anchor_grid) #The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight. + +model.to(device).eval() + +with open(wts_file, 'w') as f: + f.write('{}\n'.format(len(model.state_dict().keys()))) + for k, v in model.state_dict().items(): vr = v.reshape(-1).cpu().numpy() f.write('{} {} '.format(k, len(vr))) for vv in vr: - f.write(" ") + f.write(' ') f.write(struct.pack('>f', float(vv)).hex()) f.write('\n') diff --git a/yolov7/yololayer.cu b/yolov7/yololayer.cu index 505ec48..9e8076d 100644 --- a/yolov7/yololayer.cu +++ b/yolov7/yololayer.cu @@ -29,8 +29,8 @@ namespace nvinfer1 YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel) { mClassCount = classCount; - mYoloV5NetWidth = netWidth; - mYoloV5NetHeight = netHeight; + mYoloV7NetWidth = netWidth; + mYoloV7NetHeight = netHeight; mMaxOutObject = maxOut; /*mYoloKernel.clear(); mYoloKernel.push_back(yolo1); @@ -66,8 +66,8 @@ namespace nvinfer1 read(d, mClassCount); read(d, mThreadCount); read(d, mKernelCount); - read(d, mYoloV5NetWidth); - read(d, mYoloV5NetHeight); + read(d, mYoloV7NetWidth); + read(d, mYoloV7NetHeight); read(d, mMaxOutObject); mYoloKernel.resize(mKernelCount); auto kernelSize = mKernelCount * sizeof(YoloKernel); @@ -91,8 +91,8 @@ namespace nvinfer1 write(d, mClassCount); write(d, mThreadCount); write(d, mKernelCount); - write(d, mYoloV5NetWidth); - write(d, mYoloV5NetHeight); + write(d, mYoloV7NetWidth); + write(d, mYoloV7NetHeight); write(d, mMaxOutObject); auto kernelSize = mKernelCount * sizeof(YoloKernel); memcpy(d, mYoloKernel.data(), kernelSize); @@ -103,7 +103,7 @@ namespace nvinfer1 size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT { - return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject); + return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV7NetWidth) + sizeof(mYoloV7NetHeight) + sizeof(mMaxOutObject); } int YoloLayerPlugin::initialize() TRT_NOEXCEPT @@ -178,7 +178,7 @@ namespace nvinfer1 // Clone the plugin IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT { - YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel); + YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV7NetWidth, mYoloV7NetHeight, mMaxOutObject, mYoloKernel); p->setPluginNamespace(mPluginNamespace); return p; } @@ -260,9 +260,9 @@ namespace nvinfer1 CalDetection << < (numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > - (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem); + (inputs[i], output, numElem, mYoloV7NetWidth, mYoloV7NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem); /* CalDetection << < (numElem + mThreadCount - 1) / mThreadCount2, mThreadCount2, 0, stream >> > - (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem);*/ + (inputs[i], output, numElem, mYoloV7NetWidth, mYoloV7NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem);*/ } } diff --git a/yolov7/yololayer.h b/yolov7/yololayer.h index 87ea877..8760ee2 100644 --- a/yolov7/yololayer.h +++ b/yolov7/yololayer.h @@ -112,8 +112,8 @@ namespace nvinfer1 const char* mPluginNamespace; int mKernelCount; int mClassCount; - int mYoloV5NetWidth; - int mYoloV5NetHeight; + int mYoloV7NetWidth; + int mYoloV7NetHeight; int mMaxOutObject; std::vector mYoloKernel; void** mAnchor; diff --git a/yolov7/yolov7.cpp b/yolov7/yolov7.cpp index 47ad0e7..dc3eeab 100644 --- a/yolov7/yolov7.cpp +++ b/yolov7/yolov7.cpp @@ -1301,7 +1301,6 @@ ICudaEngine* build_engine_yolov7w6(unsigned int maxBatchSize, IBuilder* builder, - /*-----------------四个检测头对应极小，小，中，大物体检测-------------------*/ /*------------detect-----------*/ auto yolo = addYoLoLayer(network, weightMap, "model.118", std::vector{cv105_0, cv105_1, cv105_2, cv105_3}); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); @@ -1312,12 +1311,12 @@ ICudaEngine* build_engine_yolov7w6(unsigned int maxBatchSize, IBuilder* builder, #if defined(USE_FP16) config->setFlag(BuilderFlag::kFP16); #endif - /*----------------------生成engine模型-----------------------------*/ + std::cout << "Building engine, please wait for a while..." << std::endl; ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); std::cout << "Build engine successfully!" << std::endl; - /*------------------------销毁-----------------------------------*/ + network->destroy(); // Release host memory @@ -1712,7 +1711,7 @@ ICudaEngine* build_engine_yolov7(unsigned int maxBatchSize,IBuilder* builder, IB IElementWiseLayer* conv61 = convBnSilu(network, weightMap, *conv60->getOutput(0), 128, 3, 1, 1, "model.61"); ITensor* input_tensor_62[] = { conv61->getOutput(0), conv60->getOutput(0), conv59->getOutput(0), conv58->getOutput(0), conv57->getOutput(0), conv56->getOutput(0) }; IConcatenationLayer* concat62 = network->addConcatenation(input_tensor_62, 6); - concat62->setAxis(0);//提娜佳 + concat62->setAxis(0); IElementWiseLayer* conv63 = convBnSilu(network, weightMap, *concat62->getOutput(0), 256, 1, 1, 0, "model.63"); IElementWiseLayer* conv64 = convBnSilu(network, weightMap, *conv63->getOutput(0), 128, 1, 1, 0, "model.64"); @@ -1817,7 +1816,7 @@ ICudaEngine* build_engine_yolov7(unsigned int maxBatchSize,IBuilder* builder, IB ICudaEngine* build_engine_yolov7_tiny(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, std::string& wts_name) { INetworkDefinition* network = builder->createNetworkV2(0U); - // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAMEm,lkkkkk/////////////////////////////////// + ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W }); assert(data); std::map weightMap = loadWeights(wts_name); diff --git a/yolov7/yolov7_trt.py b/yolov7/yolov7_trt.py index e35cad6..4430571 100644 --- a/yolov7/yolov7_trt.py +++ b/yolov7/yolov7_trt.py @@ -34,7 +34,7 @@ def get_img_path_batches(batch_size, img_dir): def plot_one_box(x, img, color=None, label=None, line_thickness=None): """ description: Plots one bounding box on image img, - this function comes from YoLov5 project. + this function comes from YoLov7 project. param: x: a box likes [x1,y1,x2,y2] img: a opencv image object @@ -68,9 +68,9 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None): ) -class YoLov5TRT(object): +class YoLov7TRT(object): """ - description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops. + description: A YOLOv7 class that warps TensorRT ops, preprocess and postprocess ops. """ def __init__(self, engine_file_path): @@ -427,8 +427,8 @@ if __name__ == "__main__": if os.path.exists('output/'): shutil.rmtree('output/') os.makedirs('output/') - # a YoLov5TRT instance - yolov7_wrapper = YoLov5TRT(engine_file_path) + # a YoLov7TRT instance + yolov7_wrapper = YoLov7TRT(engine_file_path) try: print('batch size is', yolov7_wrapper.batch_size)