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

update gen_wts.py and some yolov5's names (#1109)

Browse Source 
* update gen_wts.py and some yolov5's names

* Update gen_wts.py
This commit is contained in:
lindsayshuo 2022-09-22 11:03:37 +08:00 committed by GitHub
parent ca5e9425f9
commit dc1e279516
6 changed files with 62 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
yolov7/README.md
View File

@ -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)

44
yolov7/gen_wts.py
View File

@ -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')

20
yolov7/yololayer.cu
View File

@ -29,8 +29,8 @@ namespace nvinfer1
YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector<Yolo::YoloKernel>& 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);*/
}
}

4
yolov7/yololayer.h
View File

@ -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<Yolo::YoloKernel> mYoloKernel;
void** mAnchor;

9
yolov7/yolov7.cpp
View File

@ -1301,7 +1301,6 @@ ICudaEngine* build_engine_yolov7w6(unsigned int maxBatchSize, IBuilder* builder,
/*-----------------四个检测头对应极小，小，中，大物体检测-------------------*/
/*------------detect-----------*/
auto yolo = addYoLoLayer(network, weightMap, "model.118", std::vector<IConvolutionLayer*>{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<std::string, Weights> weightMap = loadWeights(wts_name);

10
yolov7/yolov7_trt.py
View File

@ -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)