from __future__ import division import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import numpy as np from util import * # 输入 def get_test_input(): img = cv2.imread("dog-cycle-car.png") img = cv2.resize(img, (416,416)) # 调整输入图像维度 img_ = img[:,:,::-1].transpose((2,0,1)) # BGR -> RGB | H x W x C -> C x H x W img_ = img_[np.newaxis,:,:,:]/255.0 # 添加通道，置0，作为正则通道 img_ = torch.from_numpy(img_).float() # 转换成float img_ = Variable(img_) # 转成变量 return img_ #Takes a cfg file,returns a list of blocks. def parse_cfg(cfgfile): file = open(cfgfile, 'r') lines = file.read().split('\n') # store lines in a list lines = [x for x in lines if len(x)>0] # get rid of empty lines lines = [x for x in lines if x[0] != '#'] # get rid of comments lines = [x.rstrip().lstrip() for x in lines] # get rid of whitespaces block = {} blocks = [] for line in lines: if line[0] == "[": # a new block if len(block) != 0: # not empty blocks.append(block) # add blocks list block = {} # init blocks block["type"] = line[1:-1].rstrip() else: key, value = line.split("=") block[key.rstrip()] = value.lstrip() blocks.append(block) return blocks # 空层 class EmptyLayer(nn.Module): def __init__(self): super(EmptyLayer, self).__init__() # 调用父类方法 # 定义一个新的DetectionLayer保存检测边界框的锚点 class DetectionLayer(nn.Module): def __init__(self, anchors): super(DetectionLayer, self).__init__() self.anchors = anchors def create_modules(blocks): net_info = blocks[0] # input and pre-processing module_list = nn.ModuleList() prev_filters = 3 # depth of last conv output_filters = [] # number of output conv kernel，输出通道数量序列 for index, x in enumerate(blocks[1:]): module = nn.Sequential() # convolutional模块有卷积层、批量归一化层和leaky ReLU激活层 if (x["type"] == "convolutional"): # get layer info activation = x["activation"] try: batch_normalize = int(x["batch_normalize"]) bias = False except: batch_normalize = 0 bias = True filters = int(x["filters"]) # 卷积数量 padding = int(x["pad"]) # 填充数量 kernel_size = int(x["size"]) # 卷积核大小 stride= int(x["stride"]) # 步长 if padding: padding = (kernel_size - 1) // 2 # 运算后，宽度和高度不变 else: padding = 0 # Add conv layer conv = nn.Conv2d(prev_filters, filters, kernel_size, stride, padding, bias = bias) module.add_module("conv_{0}".format(index), conv) # Add batch norm layer if batch_normalize: bn = nn.BatchNorm2d(filters) module.add_module("batch_norm_{0}".format(index), bn) # Check activation if activation == "leaky": activn = nn.LeakyReLU(0.1, inplace = True) # 斜率0.1 module.add_module("leaky_{0}".format(index), activn) # upsample上采样层 elif (x["type"] == "upsample"): stride = int(x["stride"]) upsample = nn.Upsample(scale_factor = 2, mode = "nearest") # 或者mode="bilinear" module.add_module("upsample_module_list{}".format(index), upsample) # route路由层，路由层是获取之前层的拼接 elif (x["type"] == "route"): x["layers"] = x["layers"].split(",") # 保存start和end层号 # Start of a route start = int(x["layers"][0]) # end, if there exists one try: end = int(x["layers"][1]) except: end = 0 # 没有end # Positive anotation if start > 0: start = start - index if end > 0: end = end - index route = EmptyLayer() # 创建空层 module.add_module("route_{0}".format(index), route) if end < 0: # 计算卷积数量，即两层叠加 filters = output_filters[index + start] + output_filters[index + end] else: filters = output_filters[index + start] # shortcut捷径层（跳过连接），捷径层是将前一层的特征图添加到后面的层上 elif (x["type"] == "shortcut"): shortcut = EmptyLayer() module.add_module("shortcut_{}".format(index), shortcut) # yolo层，检测层 elif (x["type"] == "yolo"): # 保存mask序号 mask = x["mask"].split(",") mask = [int(x) for x in mask] # 保存anchors box anchors = x["anchors"].split(",") anchors = [int(a) for a in anchors] # 两个一组，还ge和宽 anchors = [(anchors[i], anchors[i+1]) for i in range(0, len(anchors), 2)] # 选取mask序号对应的anchors box，一般为3个 anchors = [anchors[i] for i in mask] detection = DetectionLayer(anchors) module.add_module("Detection_{}".format(index), detection) module_list.append(module) prev_filters = filters output_filters.append(filters) return (net_info, module_list) # 测试解析YOLO_v3配置文件 # blocks = parse_cfg("cfg/yolov3.cfg") # print(create_modules(blocks)) class Darknet(nn.Module): # 用net_info和module_list对网络进行初始化 def __init__(self, cfgfile): super(Darknet, self).__init__() self.blocks = parse_cfg(cfgfile) self.net_info, self.module_list = create_modules(self.blocks) # CUDA为true，则用GPU加速前向传播 def forward(self, x, CUDA): # delf.blocks第一个元素是net块 modules = self.blocks[1:] # 缓存每个层的输出特征图，以备route层和shortcut层使用 outputs = {} write = 0 # 是否遇到第一个检测图flag for i, module in enumerate(modules): module_type = (module["type"]) if module_type == "convolutional" or module_type == "upsample": x = self.module_list[i](x) elif module_type == "route": layers = module["layers"] layers = [int(a) for a in layers] if layers[0] > 0: layers[0] = layers[0] - i if len(layers) == 1: x = outputs[i + layers[0]] else: if layers[1] > 0: layers[1] = layers[1] - i map1 = outputs[i + layers[0]] map2 = outputs[i + layers[1]] x = torch.cat((map1, map2), 1) # 参数置1代表沿深度级联两个特征图 elif module_type == "shortcut": from_ = int(module["from"]) x = outputs[i-1] + outputs[i+from_] elif module_type == "yolo": anchors = self.module_list[i][0].anchors # input dimensions inp_dim = int(self.net_info["height"]) # number of classes num_classes = int(module["classes"]) # transform x = x.data x = predict_transform(x, inp_dim, anchors, num_classes, CUDA) if type(x) == int: continue # 如果收集器（容纳检测的张量）没有初始化 if not write: detections = x write = 1 else: detections = torch.cat((detections, x), 1) outputs[i] = x return detections # 测试向前传播 # model = Darknet("cfg/yolov3.cfg") # inp = get_test_input() # pred = model(inp) # print(pred) # 张量形状1x10647x85，第一个维度是批量大小；85行，包括4个边界框属性(bx,by,bh,bw)、1个objectness分数和80个类别分数 def load_weights(self, weightfile): """ 权重属于归一化层和卷积层，权重存储顺序与配置文件层级顺序一致。 conv有shortcut，shortcut连接另一个conv，则先包含先前conv权重。 conv with batch norm：bn biases,bn weights,bn running_mean,bn running_var,conv weights conv no batch norm：conv biases,conv weights """ fp = open(weightfile, "rb") # 标题信息：主版本，次版本，子版本，训练期间网络看到的图像 header = np.fromfile(fp, dtype = np.int32, count = 5) self.header = torch.from_numpy(header) self.seen = self.header[3] weights = np.fromfile(fp, dtype = np.float32) # 迭代地加载权重文件到网络的模块上 ptr = 0 # 追踪权重数组位置指针 for i in range(len(self.module_list)): module_type = self.blocks[i + 1]["type"] # 块包含第一块，模块不包含第一块 if module_type == "convolutional": model = self.module_list[i] # 根据conv模块是否有batch_normalize，加载权重 try: batch_normalize = int(self.blocks[i+1]["batch_normalize"]) except: batch_normalize = 0 conv = model[0] # conv with batch norm if (batch_normalize): bn = model[1] # 获取b_n layer权重的数量 num_bn_bias = bn.bias.numel() # 加载权重 bn_bias = torch.from_numpy(weights[ptr:ptr+num_bn_bias]) ptr += num_bn_bias bn_weight = torch.from_numpy(weights[ptr:ptr+num_bn_bias]) ptr += num_bn_bias bn_running_mean = torch.from_numpy(weights[ptr:ptr+num_bn_bias]) ptr += num_bn_bias bn_running_var = torch.from_numpy(weights[ptr:ptr+num_bn_bias]) ptr += num_bn_bias # 根据模型权重的维度调整重塑加载的权重 bn_bias = bn_bias.view_as(bn.bias.data) bn_weight = bn_weight.view_as(bn.weight.data) bn_running_mean = bn_running_mean.view_as(bn.running_mean) bn_running_var = bn_running_var.view_as(bn.running_var) # 将数据复制到模型中 bn.bias.data.copy_(bn_bias) bn.weight.data.copy_(bn_weight) bn.running_mean.copy_(bn_running_mean) bn.running_var.copy_(bn_running_var) # conv no batch norm，只加载卷积层的偏置项 else: # 偏置数量 num_bias = conv.bias.numel() # 加载权重 conv_bias = torch.from_numpy(weights[ptr:ptr+num_bias]) ptr += num_bias # 根据模型权重的维度调整重塑加载的权重 conv_bias = conv_bias.view_as(conv.bias.data) # 将数据复制到模型中 conv.bias.data.copy_(conv_bias) # 最后，加载共有的卷积层权重 num_weight = conv.weight.numel() conv_weight = torch.from_numpy(weights[ptr:ptr+num_weight]) ptr += num_weight conv_weight = conv_weight.view_as(conv.weight.data) conv.weight.data.copy_(conv_weight) # 测试加载预训练权重 # model = Darknet("cfg/yolov3.cfg") # model.load_weights("yolov3.weights") # inp = get_test_input() # pred = model(inp) # print(pred)