From b3c47d3feccbd409c9377a5640d9f2f1cf6d1a34 Mon Sep 17 00:00:00 2001 From: David OK Date: Tue, 19 Dec 2023 13:31:14 +0000 Subject: [PATCH] WIP: save work. --- build.py | 3 +- .../shakti/inference/yolo/darknet_model.py | 221 ++++++++++++++++++ 2 files changed, 222 insertions(+), 2 deletions(-) create mode 100644 python/oddkiva/shakti/inference/yolo/darknet_model.py diff --git a/build.py b/build.py index a190e6793..340b8fc55 100755 --- a/build.py +++ b/build.py @@ -63,11 +63,10 @@ try: import pybind11 + import sysconfig PYBIND11_DIR = pybind11.get_cmake_dir() - import distutils.sysconfig as sysconfig - PYTHON_INCLUDE_DIR = sysconfig.get_python_inc() PYTHON_LIBRARY = sysconfig.get_config_var("LIBDIR") except: diff --git a/python/oddkiva/shakti/inference/yolo/darknet_model.py b/python/oddkiva/shakti/inference/yolo/darknet_model.py new file mode 100644 index 000000000..ede9d7022 --- /dev/null +++ b/python/oddkiva/shakti/inference/yolo/darknet_model.py @@ -0,0 +1,221 @@ +import torch.nn as nn + +import oddkiva.shakti.inference.yolo.darknet_config import DarknetConfig + + +class Darknet(nn.Module): + + def __init__(self, darknet_config: DarknetConfig, inference:bool=False): + super(Darknet, self).__init__() + + self.inference = inference + self.training = not self.inference + + self.models = self.create_network(self.blocks) # merge conv, bn,leaky + self.loss = self.models[len(self.models) - 1] + + if self.blocks[(len(self.blocks) - 1)]['type'] == 'region': + self.anchors = self.loss.anchors + self.num_anchors = self.loss.num_anchors + self.anchor_step = self.loss.anchor_step + self.num_classes = self.loss.num_classes + + self.header = torch.IntTensor([0, 0, 0, 0]) + self.seen = 0 + + def print_network(self): + print_cfg(self.blocks) + + def create_network(self, blocks): + models = nn.ModuleList() + + prev_filters = 3 + out_filters = [] + prev_stride = 1 + out_strides = [] + conv_id = 0 + for block in blocks: + if block['type'] == 'net': + prev_filters = int(block['channels']) + continue + elif block['type'] == 'convolutional': + conv_id = conv_id + 1 + batch_normalize = int(block['batch_normalize']) + filters = int(block['filters']) + kernel_size = int(block['size']) + stride = int(block['stride']) + is_pad = int(block['pad']) + pad = (kernel_size - 1) // 2 if is_pad else 0 + activation = block['activation'] + model = nn.Sequential() + if batch_normalize: + model.add_module('conv{0}'.format(conv_id), + nn.Conv2d(prev_filters, filters, kernel_size, stride, pad, bias=False)) + model.add_module('bn{0}'.format(conv_id), nn.BatchNorm2d(filters)) + # model.add_module('bn{0}'.format(conv_id), BN2d(filters)) + else: + model.add_module('conv{0}'.format(conv_id), + nn.Conv2d(prev_filters, filters, kernel_size, stride, pad)) + if activation == 'leaky': + model.add_module('leaky{0}'.format(conv_id), nn.LeakyReLU(0.1, inplace=True)) + elif activation == 'relu': + model.add_module('relu{0}'.format(conv_id), nn.ReLU(inplace=True)) + elif activation == 'mish': + model.add_module('mish{0}'.format(conv_id), Mish()) + elif activation == 'linear': + model.add_module('linear{0}'.format(conv_id), nn.Identity()) + elif activation == 'logistic': + model.add_module('sigmoid{0}'.format(conv_id), nn.Sigmoid()) + else: + print("No convolutional activation named {}".format(activation)) + + prev_filters = filters + out_filters.append(prev_filters) + prev_stride = stride * prev_stride + out_strides.append(prev_stride) + models.append(model) + elif block['type'] == 'maxpool': + pool_size = int(block['size']) + stride = int(block['stride']) + if stride == 1 and pool_size % 2: + # You can use Maxpooldark instead, here is convenient to convert onnx. + # Example: [maxpool] size=3 stride=1 + model = nn.MaxPool2d(kernel_size=pool_size, stride=stride, padding=pool_size // 2) + elif stride == pool_size: + # You can use Maxpooldark instead, here is convenient to convert onnx. + # Example: [maxpool] size=2 stride=2 + model = nn.MaxPool2d(kernel_size=pool_size, stride=stride, padding=0) + else: + model = MaxPoolDark(pool_size, stride) + out_filters.append(prev_filters) + prev_stride = stride * prev_stride + out_strides.append(prev_stride) + models.append(model) + elif block['type'] == 'avgpool': + model = GlobalAvgPool2d() + out_filters.append(prev_filters) + models.append(model) + elif block['type'] == 'softmax': + model = nn.Softmax() + out_strides.append(prev_stride) + out_filters.append(prev_filters) + models.append(model) + elif block['type'] == 'cost': + if block['_type'] == 'sse': + model = nn.MSELoss(reduction='mean') + elif block['_type'] == 'L1': + model = nn.L1Loss(reduction='mean') + elif block['_type'] == 'smooth': + model = nn.SmoothL1Loss(reduction='mean') + out_filters.append(1) + out_strides.append(prev_stride) + models.append(model) + elif block['type'] == 'reorg': + stride = int(block['stride']) + prev_filters = stride * stride * prev_filters + out_filters.append(prev_filters) + prev_stride = prev_stride * stride + out_strides.append(prev_stride) + models.append(Reorg(stride)) + elif block['type'] == 'upsample': + stride = int(block['stride']) + out_filters.append(prev_filters) + prev_stride = prev_stride // stride + out_strides.append(prev_stride) + + models.append(Upsample_expand(stride)) + # models.append(Upsample_interpolate(stride)) + + elif block['type'] == 'route': + layers = block['layers'].split(',') + ind = len(models) + layers = [int(i) if int(i) > 0 else int(i) + ind for i in layers] + if len(layers) == 1: + if 'groups' not in block.keys() or int(block['groups']) == 1: + prev_filters = out_filters[layers[0]] + prev_stride = out_strides[layers[0]] + else: + prev_filters = out_filters[layers[0]] // int(block['groups']) + prev_stride = out_strides[layers[0]] // int(block['groups']) + elif len(layers) == 2: + assert (layers[0] == ind - 1 or layers[1] == ind - 1) + prev_filters = out_filters[layers[0]] + out_filters[layers[1]] + prev_stride = out_strides[layers[0]] + elif len(layers) == 4: + assert (layers[0] == ind - 1) + prev_filters = out_filters[layers[0]] + out_filters[layers[1]] + out_filters[layers[2]] + \ + out_filters[layers[3]] + prev_stride = out_strides[layers[0]] + else: + print("route error!!!") + + out_filters.append(prev_filters) + out_strides.append(prev_stride) + models.append(EmptyModule()) + elif block['type'] == 'shortcut': + ind = len(models) + prev_filters = out_filters[ind - 1] + out_filters.append(prev_filters) + prev_stride = out_strides[ind - 1] + out_strides.append(prev_stride) + models.append(EmptyModule()) + elif block['type'] == 'sam': + ind = len(models) + prev_filters = out_filters[ind - 1] + out_filters.append(prev_filters) + prev_stride = out_strides[ind - 1] + out_strides.append(prev_stride) + models.append(EmptyModule()) + elif block['type'] == 'connected': + filters = int(block['output']) + if block['activation'] == 'linear': + model = nn.Linear(prev_filters, filters) + elif block['activation'] == 'leaky': + model = nn.Sequential( + nn.Linear(prev_filters, filters), + nn.LeakyReLU(0.1, inplace=True)) + elif block['activation'] == 'relu': + model = nn.Sequential( + nn.Linear(prev_filters, filters), + nn.ReLU(inplace=True)) + prev_filters = filters + out_filters.append(prev_filters) + out_strides.append(prev_stride) + models.append(model) + elif block['type'] == 'region': + loss = RegionLoss() + anchors = block['anchors'].split(',') + loss.anchors = [float(i) for i in anchors] + loss.num_classes = int(block['classes']) + loss.num_anchors = int(block['num']) + loss.anchor_step = len(loss.anchors) // loss.num_anchors + loss.object_scale = float(block['object_scale']) + loss.noobject_scale = float(block['noobject_scale']) + loss.class_scale = float(block['class_scale']) + loss.coord_scale = float(block['coord_scale']) + out_filters.append(prev_filters) + out_strides.append(prev_stride) + models.append(loss) + elif block['type'] == 'yolo': + yolo_layer = YoloLayer() + anchors = block['anchors'].split(',') + anchor_mask = block['mask'].split(',') + yolo_layer.anchor_mask = [int(i) for i in anchor_mask] + yolo_layer.anchors = [float(i) for i in anchors] + yolo_layer.num_classes = int(block['classes']) + self.num_classes = yolo_layer.num_classes + yolo_layer.num_anchors = int(block['num']) + yolo_layer.anchor_step = len(yolo_layer.anchors) // yolo_layer.num_anchors + yolo_layer.stride = prev_stride + yolo_layer.scale_x_y = float(block['scale_x_y']) + # yolo_layer.object_scale = float(block['object_scale']) + # yolo_layer.noobject_scale = float(block['noobject_scale']) + # yolo_layer.class_scale = float(block['class_scale']) + # yolo_layer.coord_scale = float(block['coord_scale']) + out_filters.append(prev_filters) + out_strides.append(prev_stride) + models.append(yolo_layer) + else: + print('unknown type %s' % (block['type'])) + + return models