thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Minor fixes

Browse Source
This commit is contained in:
Marcos Luciano
2022-04-08 13:31:36 -03:00
parent a310f90cd7
commit 2aa52a8e8c
3 changed files with 121 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -174,9 +174,9 @@ NvDsInferStatus Yolo::buildYoloNetwork(
else if (m_ConfigBlocks.at(i).at("type") == "convolutional") { else if (m_ConfigBlocks.at(i).at("type") == "convolutional") {
float eps = 1.0e-5; float eps = 1.0e-5;
if (m_NetworkType.find("yolov5") != std::string::npos) { /*if (m_NetworkType.find("yolov5") != std::string::npos) {
eps = 1.0e-3; eps = 1.0e-3;
} }*/
std::string inputVol = dimsToString(previous->getDimensions()); std::string inputVol = dimsToString(previous->getDimensions());
nvinfer1::ILayer* out = convolutionalLayer(i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, weightsType, channels, eps, previous, &network); nvinfer1::ILayer* out = convolutionalLayer(i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, weightsType, channels, eps, previous, &network);
previous = out->getOutput(0); previous = out->getOutput(0);
@@ -233,11 +233,12 @@ NvDsInferStatus Yolo::buildYoloNetwork(
printLayerInfo(layerIndex, layerType, " -", outputVol, " -"); printLayerInfo(layerIndex, layerType, " -", outputVol, " -");
} }
else if (m_ConfigBlocks.at(i).at("type") == "dropout") { // Skip dropout layer else if (m_ConfigBlocks.at(i).at("type") == "dropout") {
// Skip dropout layer
assert(m_ConfigBlocks.at(i).find("probability") != m_ConfigBlocks.at(i).end()); assert(m_ConfigBlocks.at(i).find("probability") != m_ConfigBlocks.at(i).end());
//float probability = std::stof(m_ConfigBlocks.at(i).at("probability")); /*float probability = std::stof(m_ConfigBlocks.at(i).at("probability"));
//nvinfer1::ILayer* out = dropoutLayer(probability, previous, &network); nvinfer1::ILayer* out = dropoutLayer(probability, previous, &network);
//previous = out->getOutput(0); previous = out->getOutput(0);*/
assert(previous != nullptr); assert(previous != nullptr);
tensorOutputs.push_back(previous); tensorOutputs.push_back(previous);
printLayerInfo(layerIndex, "dropout", " -", " -", " -"); printLayerInfo(layerIndex, "dropout", " -", " -", " -");

277
utils/gen_wts_yoloV5.py
View File

@@ -7,6 +7,57 @@ import torch
from utils.torch_utils import select_device from utils.torch_utils import select_device
class YoloLayers():
def get_route(self, n, layers):
route = 0
for i, layer in enumerate(layers):
if i <= n:
route += layer[1]
else:
break
return route
def route(self, layers=""):
return "\n[route]\n" + \
"layers=%s\n" % layers
def shortcut(self, route=-1, activation="linear"):
return "\n[shortcut]\n" + \
"from=%d\n" % route + \
"activation=%s\n" % activation
def maxpool(self, stride=1, size=1):
return "\n[maxpool]\n" + \
"stride=%d\n" % stride + \
"size=%d\n" % size
def upsample(self, stride=1):
return "\n[upsample]\n" + \
"stride=%d\n" % stride
def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation="linear"):
b = "batch_normalize=1\n" if bn is True else ""
g = "groups=%d\n" % groups if groups > 1 else ""
return "\n[convolutional]\n" + \
b + \
"filters=%d\n" % filters + \
"size=%d\n" % size + \
"stride=%d\n" % stride + \
"pad=%d\n" % pad + \
g + \
"activation=%s\n" % activation
def yolo(self, mask="", anchors="", classes=80, num=3):
return "\n[yolo]\n" + \
"mask=%s\n" % mask + \
"anchors=%s\n" % anchors + \
"classes=%d\n" % classes + \
"num=%d\n" % num + \
"scale_x_y=2.0\n" + \
"beta_nms=0.6\n" + \
"new_coords=1\n"
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description="PyTorch YOLOv5 conversion") parser = argparse.ArgumentParser(description="PyTorch YOLOv5 conversion")
parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)") parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)")
@@ -77,10 +128,10 @@ with open(wts_file, "w") as f:
cv3_idx = 0 cv3_idx = 0
sppf_idx = 11 if p6 else 9 sppf_idx = 11 if p6 else 9
for k, v in model.state_dict().items(): for k, v in model.state_dict().items():
if not "num_batches_tracked" in k and not "anchors" in k and not "anchor_grid" in k: if "num_batches_tracked" not in k and "anchors" not in k and "anchor_grid" not in k:
vr = v.reshape(-1).cpu().numpy() vr = v.reshape(-1).cpu().numpy()
idx = int(k.split(".")[1]) idx = int(k.split(".")[1])
if ".cv1." in k and not ".m." in k and idx != sppf_idx: if ".cv1." in k and ".m." not in k and idx != sppf_idx:
cv1 += "{} {} ".format(k, len(vr)) cv1 += "{} {} ".format(k, len(vr))
for vv in vr: for vv in vr:
cv1 += " " cv1 += " "
@@ -102,7 +153,7 @@ with open(wts_file, "w") as f:
wts_write += cv3 wts_write += cv3
cv3 = "" cv3 = ""
cv3_idx = 0 cv3_idx = 0
if not ".cv3." in k and not (".cv1." in k and not ".m." in k and idx != sppf_idx): if ".cv3." not in k and not (".cv1." in k and ".m." not in k and idx != sppf_idx):
wts_write += "{} {} ".format(k, len(vr)) wts_write += "{} {} ".format(k, len(vr))
for vv in vr: for vv in vr:
wts_write += " " wts_write += " "
@@ -125,219 +176,117 @@ with open(wts_file, "w") as f:
with open(cfg_file, "w") as c: with open(cfg_file, "w") as c:
with open(yaml_file, "r", encoding="utf-8") as f: with open(yaml_file, "r", encoding="utf-8") as f:
nc = 0
depth_multiple = 0
width_multiple = 0
detections = []
layers = []
f = yaml.load(f,Loader=yaml.FullLoader)
c.write("[net]\n") c.write("[net]\n")
c.write("width=%d\n" % model_width) c.write("width=%d\n" % model_width)
c.write("height=%d\n" % model_height) c.write("height=%d\n" % model_height)
c.write("channels=%d\n" % model_channels) c.write("channels=%d\n" % model_channels)
for l in f: nc = 0
if l == "nc": depth_multiple = 0
nc = f[l] width_multiple = 0
elif l == "depth_multiple": layers = []
depth_multiple = f[l] yoloLayers = YoloLayers()
elif l == "width_multiple": f = yaml.load(f, Loader=yaml.FullLoader)
width_multiple = f[l] for topic in f:
elif l == "backbone" or l == "head": if topic == "nc":
for v in f[l]: nc = f[topic]
elif topic == "depth_multiple":
depth_multiple = f[topic]
elif topic == "width_multiple":
width_multiple = f[topic]
elif topic == "backbone" or topic == "head":
for v in f[topic]:
if v[2] == "Conv": if v[2] == "Conv":
layer = "" layer = "\n# Conv\n"
blocks = 0 blocks = 0
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
layer += "batch_normalize=1\n" stride=v[3][2], activation="silu")
layer += "filters=%d\n" % get_width(v[3][0], width_multiple)
layer += "size=%d\n" % v[3][1]
layer += "stride=%d\n" % v[3][2]
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "C3": elif v[2] == "C3":
layer = "" layer = "\n# C3\n"
blocks = 0 blocks = 0
layer += "\n# C3\n"
# SPLIT # SPLIT
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "batch_normalize=1\n" activation="silu")
layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
layer += "\n[route]\n" layer += yoloLayers.route(layers="-2")
layer += "layers=-2\n"
blocks += 1 blocks += 1
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "batch_normalize=1\n" activation="silu")
layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
# Residual Block # Residual Block
if len(v[3]) == 1 or v[3][1] == True: if len(v[3]) == 1 or v[3][1] is True:
for _ in range(get_depth(v[1], depth_multiple)): for _ in range(get_depth(v[1], depth_multiple)):
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "batch_normalize=1\n" activation="silu")
layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "batch_normalize=1\n" size=3, activation="silu")
layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple)
layer += "size=3\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
layer += "\n[shortcut]\n" layer += yoloLayers.shortcut(route=-3)
layer += "from=-3\n"
layer += "activation=linear\n"
blocks += 1 blocks += 1
# Merge # Merge
layer += "\n[route]\n" layer += yoloLayers.route(layers="-1, -%d" % (3 * get_depth(v[1], depth_multiple) + 3))
layer += "layers=-1, -%d\n" % (3 * get_depth(v[1], depth_multiple) + 3)
blocks += 1 blocks += 1
else: else:
for _ in range(get_depth(v[1], depth_multiple)): for _ in range(get_depth(v[1], depth_multiple)):
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "batch_normalize=1\n" activation="silu")
layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "batch_normalize=1\n" size=3, activation="silu")
layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple)
layer += "size=3\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
# Merge # Merge
layer += "\n[route]\n" layer += yoloLayers.route(layers="-1, -%d" % (2 * get_depth(v[1], depth_multiple) + 3))
layer += "layers=-1, -%d\n" % (2 * get_depth(v[1], depth_multiple) + 3)
blocks += 1 blocks += 1
# Transition # Transition
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
layer += "batch_normalize=1\n" activation="silu")
layer += "filters=%d\n" % get_width(v[3][0], width_multiple)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
layer += "\n##########\n"
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "SPPF": elif v[2] == "SPPF":
layer = "" layer = "\n# SPPF\n"
blocks = 0 blocks = 0
layer += "\n# SPPF\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
layer += "\n[convolutional]\n" activation="silu")
layer += "batch_normalize=1\n"
layer += "filters=%d\n" % (get_width(v[3][0], width_multiple) / 2)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
blocks += 1 blocks += 1
layer += "\n[maxpool]\n" layer += yoloLayers.maxpool(size=v[3][1])
layer += "stride=1\n"
layer += "size=%d\n" % v[3][1]
blocks += 1 blocks += 1
layer += "\n[maxpool]\n" layer += yoloLayers.maxpool(size=v[3][1])
layer += "stride=1\n"
layer += "size=%d\n" % v[3][1]
blocks += 1 blocks += 1
layer += "\n[maxpool]\n" layer += yoloLayers.maxpool(size=v[3][1])
layer += "stride=1\n"
layer += "size=%d\n" % v[3][1]
blocks += 1 blocks += 1
layer += "\n[route]\n" layer += yoloLayers.route(layers="-4, -3, -2, -1")
layer += "layers=-4, -3, -2, -1\n"
blocks += 1 blocks += 1
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
layer += "batch_normalize=1\n" activation="silu")
layer += "filters=%d\n" % get_width(v[3][0], width_multiple)
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "activation=silu\n"
layer += "\n##########\n"
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "nn.Upsample": elif v[2] == "nn.Upsample":
layer = "" layer = "\n# nn.Upsample\n"
blocks = 0 blocks = 0
layer += "\n[upsample]\n" layer += yoloLayers.upsample(stride=v[3][1])
layer += "stride=%d\n" % v[3][1]
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "Concat": elif v[2] == "Concat":
layer = ""
blocks = 0
route = v[0][1] route = v[0][1]
r = 0 route = yoloLayers.get_route(route, layers) if route > 0 else \
if route > 0: yoloLayers.get_route(len(layers) + route, layers)
for i, item in enumerate(layers): layer = "\n# Concat\n"
if i <= route: blocks = 0
r += item[1] layer += yoloLayers.route(layers="-1, %d" % (route - 1))
else:
break
else:
route = len(layers) + route
for i, item in enumerate(layers):
if i <= route:
r += item[1]
else:
break
layer += "\n# Concat\n"
layer += "\n[route]\n"
layer += "layers=-1, %d\n" % (r - 1)
layer += "\n##########\n"
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
elif v[2] == "Detect": elif v[2] == "Detect":
for i, n in enumerate(v[0]): for i, n in enumerate(v[0]):
layer = "" route = yoloLayers.get_route(n, layers)
layer = "\n# Detect\n"
blocks = 0 blocks = 0
r = 0 layer += yoloLayers.route(layers="%d" % (route - 1))
for j, item in enumerate(layers):
if j <= n:
r += item[1]
else:
break
layer += "\n# Detect\n"
layer += "\n[route]\n"
layer += "layers=%d\n" % (r - 1)
blocks += 1 blocks += 1
layer += "\n[convolutional]\n" layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation="logistic")
layer += "size=1\n"
layer += "stride=1\n"
layer += "pad=1\n"
layer += "filters=%d\n" % ((nc + 5) * len(masks[i]))
layer += "activation=logistic\n"
blocks += 1 blocks += 1
layer += "\n[yolo]\n" layer += yoloLayers.yolo(mask=str(masks[i])[1:-1], anchors=anchors, classes=nc, num=num)
layer += "mask=%s\n" % str(masks[i])[1:-1]
layer += "anchors=%s\n" % anchors
layer += "classes=%d\n" % nc
layer += "num=%d\n" % num
layer += "scale_x_y=2.0\n"
layer += "beta_nms=0.6\n"
layer += "new_coords=1\n"
layer += "\n##########\n"
blocks += 1 blocks += 1
layers.append([layer, blocks]) layers.append([layer, blocks])
for layer in layers: for layer in layers:

2
utils/gen_wts_yolor.py
View File

@@ -31,7 +31,7 @@ with open(wts_file, "w") as f:
wts_write = "" wts_write = ""
conv_count = 0 conv_count = 0
for k, v in model.state_dict().items(): for k, v in model.state_dict().items():
if not "num_batches_tracked" in k: if "num_batches_tracked" not in k:
vr = v.reshape(-1).cpu().numpy() vr = v.reshape(-1).cpu().numpy()
wts_write += "{} {} ".format(k, len(vr)) wts_write += "{} {} ".format(k, len(vr))
for vv in vr: for vv in vr: