import torch.nn as nn import torch.nn.functional as F import numpy as np from tool.region_loss import RegionLoss from tool.yolo_layer import YoloLayer from tool.config import * from tool.torch_utils import * class Mish(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): x = x * (torch.tanh(torch.nn.functional.softplus(x))) return x class MaxPoolDark(nn.Module): def __init__(self, size=2, stride=1): super(MaxPoolDark, self).__init__() self.size = size self.stride = stride def forward(self, x): ''' darknet output_size = (input_size + p - k) / s +1 p : padding = k - 1 k : size s : stride torch output_size = (input_size + 2*p -k) / s +1 p : padding = k//2 ''' p = self.size // 2 if ((x.shape[2] - 1) // self.stride) != ((x.shape[2] + 2 * p - self.size) // self.stride): padding1 = (self.size - 1) // 2 padding2 = padding1 + 1 else: padding1 = (self.size - 1) // 2 padding2 = padding1 if ((x.shape[3] - 1) // self.stride) != ((x.shape[3] + 2 * p - self.size) // self.stride): padding3 = (self.size - 1) // 2 padding4 = padding3 + 1 else: padding3 = (self.size - 1) // 2 padding4 = padding3 x = F.max_pool2d(F.pad(x, (padding3, padding4, padding1, padding2), mode='replicate'), self.size, stride=self.stride) return x class Upsample_expand(nn.Module): def __init__(self, stride=2): super(Upsample_expand, self).__init__() self.stride = stride def forward(self, x): assert (x.data.dim() == 4) x = x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\ expand(x.size(0), x.size(1), x.size(2), self.stride, x.size(3), self.stride).contiguous().\ view(x.size(0), x.size(1), x.size(2) * self.stride, x.size(3) * self.stride) return x class Upsample_interpolate(nn.Module): def __init__(self, stride): super(Upsample_interpolate, self).__init__() self.stride = stride def forward(self, x): assert (x.data.dim() == 4) out = F.interpolate(x, size=(x.size(2) * self.stride, x.size(3) * self.stride), mode='nearest') return out class Reorg(nn.Module): def __init__(self, stride=2): super(Reorg, self).__init__() self.stride = stride def forward(self, x): stride = self.stride assert (x.data.dim() == 4) B = x.data.size(0) C = x.data.size(1) H = x.data.size(2) W = x.data.size(3) assert (H % stride == 0) assert (W % stride == 0) ws = stride hs = stride x = x.view(B, C, H / hs, hs, W / ws, ws).transpose(3, 4).contiguous() x = x.view(B, C, H / hs * W / ws, hs * ws).transpose(2, 3).contiguous() x = x.view(B, C, hs * ws, H / hs, W / ws).transpose(1, 2).contiguous() x = x.view(B, hs * ws * C, H / hs, W / ws) return x class GlobalAvgPool2d(nn.Module): def __init__(self): super(GlobalAvgPool2d, self).__init__() def forward(self, x): N = x.data.size(0) C = x.data.size(1) H = x.data.size(2) W = x.data.size(3) x = F.avg_pool2d(x, (H, W)) x = x.view(N, C) return x # for route, shortcut and sam class EmptyModule(nn.Module): def __init__(self): super(EmptyModule, self).__init__() def forward(self, x): return x # support route shortcut and reorg class Darknet(nn.Module): def __init__(self, cfgfile, inference=False): super(Darknet, self).__init__() self.inference = inference self.training = not self.inference self.blocks = parse_cfg(cfgfile) self.width = int(self.blocks[0]['width']) self.height = int(self.blocks[0]['height']) 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 forward(self, x): ind = -2 self.loss = None outputs = dict() out_boxes = [] for block in self.blocks: ind = ind + 1 # if ind > 0: # return x if block['type'] == 'net': continue elif block['type'] in ['convolutional', 'maxpool', 'reorg', 'upsample', 'avgpool', 'softmax', 'connected']: x = self.models[ind](x) outputs[ind] = x elif block['type'] == 'route': layers = block['layers'].split(',') 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: x = outputs[layers[0]] outputs[ind] = x else: groups = int(block['groups']) group_id = int(block['group_id']) _, b, _, _ = outputs[layers[0]].shape x = outputs[layers[0]][:, b // groups * group_id:b // groups * (group_id + 1)] outputs[ind] = x elif len(layers) == 2: x1 = outputs[layers[0]] x2 = outputs[layers[1]] x = torch.cat((x1, x2), 1) outputs[ind] = x elif len(layers) == 4: x1 = outputs[layers[0]] x2 = outputs[layers[1]] x3 = outputs[layers[2]] x4 = outputs[layers[3]] x = torch.cat((x1, x2, x3, x4), 1) outputs[ind] = x else: print("rounte number > 2 ,is {}".format(len(layers))) elif block['type'] == 'shortcut': from_layer = int(block['from']) activation = block['activation'] from_layer = from_layer if from_layer > 0 else from_layer + ind x1 = outputs[from_layer] x2 = outputs[ind - 1] x = x1 + x2 if activation == 'leaky': x = F.leaky_relu(x, 0.1, inplace=True) elif activation == 'relu': x = F.relu(x, inplace=True) outputs[ind] = x elif block['type'] == 'sam': from_layer = int(block['from']) from_layer = from_layer if from_layer > 0 else from_layer + ind x1 = outputs[from_layer] x2 = outputs[ind - 1] x = x1 * x2 outputs[ind] = x elif block['type'] == 'region': continue if self.loss: self.loss = self.loss + self.models[ind](x) else: self.loss = self.models[ind](x) outputs[ind] = None elif block['type'] == 'yolo': # if self.training: # pass # else: # boxes = self.models[ind](x) # out_boxes.append(boxes) boxes = self.models[ind](x) out_boxes.append(boxes) elif block['type'] == 'cost': continue else: print('unknown type %s' % (block['type'])) if self.training: return out_boxes else: return get_region_boxes(out_boxes) 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 def load_weights(self, weightfile): fp = open(weightfile, 'rb') header = np.fromfile(fp, count=5, dtype=np.int32) self.header = torch.from_numpy(header) self.seen = self.header[3] buf = np.fromfile(fp, dtype=np.float32) fp.close() start = 0 ind = -2 for block in self.blocks: if start >= buf.size: break ind = ind + 1 if block['type'] == 'net': continue elif block['type'] == 'convolutional': model = self.models[ind] batch_normalize = int(block['batch_normalize']) if batch_normalize: start = load_conv_bn(buf, start, model[0], model[1]) else: start = load_conv(buf, start, model[0]) elif block['type'] == 'connected': model = self.models[ind] if block['activation'] != 'linear': start = load_fc(buf, start, model[0]) else: start = load_fc(buf, start, model) elif block['type'] == 'maxpool': pass elif block['type'] == 'reorg': pass elif block['type'] == 'upsample': pass elif block['type'] == 'route': pass elif block['type'] == 'shortcut': pass elif block['type'] == 'sam': pass elif block['type'] == 'region': pass elif block['type'] == 'yolo': pass elif block['type'] == 'avgpool': pass elif block['type'] == 'softmax': pass elif block['type'] == 'cost': pass else: print('unknown type %s' % (block['type'])) # def save_weights(self, outfile, cutoff=0): # if cutoff <= 0: # cutoff = len(self.blocks) - 1 # # fp = open(outfile, 'wb') # self.header[3] = self.seen # header = self.header # header.numpy().tofile(fp) # # ind = -1 # for blockId in range(1, cutoff + 1): # ind = ind + 1 # block = self.blocks[blockId] # if block['type'] == 'convolutional': # model = self.models[ind] # batch_normalize = int(block['batch_normalize']) # if batch_normalize: # save_conv_bn(fp, model[0], model[1]) # else: # save_conv(fp, model[0]) # elif block['type'] == 'connected': # model = self.models[ind] # if block['activation'] != 'linear': # save_fc(fc, model) # else: # save_fc(fc, model[0]) # elif block['type'] == 'maxpool': # pass # elif block['type'] == 'reorg': # pass # elif block['type'] == 'upsample': # pass # elif block['type'] == 'route': # pass # elif block['type'] == 'shortcut': # pass # elif block['type'] == 'sam': # pass # elif block['type'] == 'region': # pass # elif block['type'] == 'yolo': # pass # elif block['type'] == 'avgpool': # pass # elif block['type'] == 'softmax': # pass # elif block['type'] == 'cost': # pass # else: # print('unknown type %s' % (block['type'])) # fp.close()