睿智的目標檢測26——Pytorch搭建yolo3目標檢測平臺

學習前言

一起來看看yolo3的Pytorch實現吧,順便訓練一下自己的數據。
在這裏插入圖片描述

源碼下載

https://github.com/bubbliiiing/yolo3-pytorch
喜歡的可以點個star噢。

yolo3實現思路

一、預測部分

1、主題網絡darknet53介紹

在這裏插入圖片描述
YOLOv3相比於之前的yolo1和yolo2,改進較大,主要改進方向有:

1、主幹網絡修改爲darknet53,其重要特點是使用了殘差網絡Residual,darknet53中的殘差卷積就是進行一次3X3、步長爲2的卷積,然後保存該卷積layer,再進行一次1X1的卷積和一次3X3的卷積,並把這個結果加上layer作爲最後的結果, 殘差網絡的特點是容易優化,並且能夠通過增加相當的深度來提高準確率。其內部的殘差塊使用了跳躍連接,緩解了在深度神經網絡中增加深度帶來的梯度消失問題。

2、darknet53的每一個卷積部分使用了特有的DarknetConv2D結構,每一次卷積的時候進行l2正則化,完成卷積後進行BatchNormalization標準化與LeakyReLU。普通的ReLU是將所有的負值都設爲零,Leaky ReLU則是給所有負值賦予一個非零斜率。以數學的方式我們可以表示爲
在這裏插入圖片描述
實現代碼爲:

import torch
import torch.nn as nn
import math
from collections import OrderedDict

# 基本的darknet塊
class BasicBlock(nn.Module):
    def __init__(self, inplanes, planes):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1,
                               stride=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(planes[0])
        self.relu1 = nn.LeakyReLU(0.1)
        
        self.conv2 = nn.Conv2d(planes[0], planes[1], kernel_size=3,
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes[1])
        self.relu2 = nn.LeakyReLU(0.1)

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu1(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu2(out)

        out += residual
        return out


class DarkNet(nn.Module):
    def __init__(self, layers):
        super(DarkNet, self).__init__()
        self.inplanes = 32
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(self.inplanes)
        self.relu1 = nn.LeakyReLU(0.1)

        self.layer1 = self._make_layer([32, 64], layers[0])
        self.layer2 = self._make_layer([64, 128], layers[1])
        self.layer3 = self._make_layer([128, 256], layers[2])
        self.layer4 = self._make_layer([256, 512], layers[3])
        self.layer5 = self._make_layer([512, 1024], layers[4])

        self.layers_out_filters = [64, 128, 256, 512, 1024]

        # 進行權值初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, planes, blocks):
        layers = []
        # 下采樣,步長爲2,卷積核大小爲3
        layers.append(("ds_conv", nn.Conv2d(self.inplanes, planes[1], kernel_size=3,
                                stride=2, padding=1, bias=False)))
        layers.append(("ds_bn", nn.BatchNorm2d(planes[1])))
        layers.append(("ds_relu", nn.LeakyReLU(0.1)))
        # 加入darknet模塊   
        self.inplanes = planes[1]
        for i in range(0, blocks):
            layers.append(("residual_{}".format(i), BasicBlock(self.inplanes, planes)))
        return nn.Sequential(OrderedDict(layers))

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)

        x = self.layer1(x)
        x = self.layer2(x)
        out3 = self.layer3(x)
        out4 = self.layer4(out3)
        out5 = self.layer5(out4)

        return out3, out4, out5

def darknet53(pretrained, **kwargs):
    model = DarkNet([1, 2, 8, 8, 4])
    if pretrained:
        if isinstance(pretrained, str):
            model.load_state_dict(torch.load(pretrained))
        else:
            raise Exception("darknet request a pretrained path. got [{}]".format(pretrained))
    return model

2、從特徵獲取預測結果

在這裏插入圖片描述
1、在特徵利用部分,yolo3提取多特徵層進行目標檢測,一共提取三個特徵層,三個特徵層位於主幹部分darknet53的不同位置,分別位於中間層,中下層,底層,三個特徵層的shape分別爲(52,52,256)、(26,26,512)、(13,13,1024)。

2、三個特徵層進行5次卷積處理,處理完後一部分用於輸出該特徵層對應的預測結果一部分用於進行反捲積UmSampling2d後與其它特徵層進行結合

3、輸出層的shape分別爲(13,13,75),(26,26,75),(52,52,75),最後一個維度爲75是因爲該圖是基於voc數據集的,它的類爲20種,yolo3只有針對每一個特徵層存在3個先驗框,所以最後維度爲3x25;
如果使用的是coco訓練集,類則爲80種,最後的維度應該爲255 = 3x85
,三個特徵層的shape爲(13,13,255),(26,26,255),(52,52,255)

其實際情況就是,由於我們使用得是Pytorch,它的通道數默認在第一位,輸入N張416x416的圖片,在經過多層的運算後,會輸出三個shape分別爲(N,255,13,13),(N,255,26,26),(N,255,52,52)的數據,對應每個圖分爲13x13、26x26、52x52的網格上3個先驗框的位置。
實現代碼如下:

import torch
import torch.nn as nn
from collections import OrderedDict
from nets.darknet import darknet53

def conv2d(filter_in, filter_out, kernel_size):
    pad = (kernel_size - 1) // 2 if kernel_size else 0
    return nn.Sequential(OrderedDict([
        ("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=1, padding=pad, bias=False)),
        ("bn", nn.BatchNorm2d(filter_out)),
        ("relu", nn.LeakyReLU(0.1)),
    ]))

def make_last_layers(filters_list, in_filters, out_filter):
    m = nn.ModuleList([
        conv2d(in_filters, filters_list[0], 1),
        conv2d(filters_list[0], filters_list[1], 3),
        conv2d(filters_list[1], filters_list[0], 1),
        conv2d(filters_list[0], filters_list[1], 3),
        conv2d(filters_list[1], filters_list[0], 1),
        conv2d(filters_list[0], filters_list[1], 3),
        nn.Conv2d(filters_list[1], out_filter, kernel_size=1,
                                        stride=1, padding=0, bias=True)
    ])
    return m

class YoloBody(nn.Module):
    def __init__(self, config):
        super(YoloBody, self).__init__()
        self.config = config
        #  backbone
        self.backbone = darknet53(None)

        out_filters = self.backbone.layers_out_filters
        #  last_layer0
        final_out_filter0 = len(config["yolo"]["anchors"][0]) * (5 + config["yolo"]["classes"])
        self.last_layer0 = make_last_layers([512, 1024], out_filters[-1], final_out_filter0)

        #  embedding1
        final_out_filter1 = len(config["yolo"]["anchors"][1]) * (5 + config["yolo"]["classes"])
        self.last_layer1_conv = conv2d(512, 256, 1)
        self.last_layer1_upsample = nn.Upsample(scale_factor=2, mode='nearest')
        self.last_layer1 = make_last_layers([256, 512], out_filters[-2] + 256, final_out_filter1)

        #  embedding2
        final_out_filter2 = len(config["yolo"]["anchors"][2]) * (5 + config["yolo"]["classes"])
        self.last_layer2_conv = conv2d(256, 128, 1)
        self.last_layer2_upsample = nn.Upsample(scale_factor=2, mode='nearest')
        self.last_layer2 = make_last_layers([128, 256], out_filters[-3] + 128, final_out_filter2)


    def forward(self, x):
        def _branch(last_layer, layer_in):
            for i, e in enumerate(last_layer):
                layer_in = e(layer_in)
                if i == 4:
                    out_branch = layer_in
            return layer_in, out_branch
        #  backbone
        x2, x1, x0 = self.backbone(x)
        #  yolo branch 0
        out0, out0_branch = _branch(self.last_layer0, x0)

        #  yolo branch 1
        x1_in = self.last_layer1_conv(out0_branch)
        x1_in = self.last_layer1_upsample(x1_in)
        x1_in = torch.cat([x1_in, x1], 1)
        out1, out1_branch = _branch(self.last_layer1, x1_in)

        #  yolo branch 2
        x2_in = self.last_layer2_conv(out1_branch)
        x2_in = self.last_layer2_upsample(x2_in)
        x2_in = torch.cat([x2_in, x2], 1)
        out2, _ = _branch(self.last_layer2, x2_in)
        return out0, out1, out2

3、預測結果的解碼

由第二步我們可以獲得三個特徵層的預測結果,shape分別爲(N,255,13,13),(N,255,26,26),(N,255,52,52)的數據,對應每個圖分爲13x13、26x26、52x52的網格上3個預測框的位置。

但是這個預測結果並不對應着最終的預測框在圖片上的位置,還需要解碼纔可以完成。

此處要講一下yolo3的預測原理,yolo3的3個特徵層分別將整幅圖分爲13x13、26x26、52x52的網格,每個網絡點負責一個區域的檢測。

我們知道特徵層的預測結果對應着三個預測框的位置,我們先將其reshape一下,其結果爲(N,3,85,13,13,3,85),(N,3,85,26,26),(N,3,85,52,52)。

維度中的85包含了4+1+80,分別代表x_offset、y_offset、h和w、置信度、分類結果。

yolo3的解碼過程就是將每個網格點加上它對應的x_offset和y_offset,加完後的結果就是預測框的中心,然後再利用 先驗框和h、w結合 計算出預測框的長和寬。這樣就能得到整個預測框的位置了。

在這裏插入圖片描述
當然得到最終的預測結構後還要進行得分排序與非極大抑制篩選
這一部分基本上是所有目標檢測通用的部分。不過該項目的處理方式與其它項目不同。其對於每一個類進行判別。
1、取出每一類得分大於self.obj_threshold的框和得分。
2、利用框的位置和得分進行非極大抑制。

實現代碼如下

class DecodeBox(nn.Module):
    def __init__(self, anchors, num_classes, img_size):
        super(DecodeBox, self).__init__()
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.bbox_attrs = 5 + num_classes
        self.img_size = img_size

    def forward(self, input):
        batch_size = input.size(0)
        input_height = input.size(2)
        input_width = input.size(3)

        # 計算步長
        stride_h = self.img_size[1] / input_height
        stride_w = self.img_size[0] / input_width
        # 歸一到特徵層上
        scaled_anchors = [(anchor_width / stride_w, anchor_height / stride_h) for anchor_width, anchor_height in self.anchors]

        # 對預測結果進行resize
        prediction = input.view(batch_size, self.num_anchors,
                                self.bbox_attrs, input_height, input_width).permute(0, 1, 3, 4, 2).contiguous()

        # 先驗框的中心位置的調整參數
        x = torch.sigmoid(prediction[..., 0])  
        y = torch.sigmoid(prediction[..., 1])
        # 先驗框的寬高調整參數
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height

        # 獲得置信度,是否有物體
        conf = torch.sigmoid(prediction[..., 4])
        # 種類置信度
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor

        # 生成網格,先驗框中心,網格左上角
        grid_x = torch.linspace(0, input_width - 1, input_width).repeat(input_width, 1).repeat(
            batch_size * self.num_anchors, 1, 1).view(x.shape).type(FloatTensor)
        grid_y = torch.linspace(0, input_height - 1, input_height).repeat(input_height, 1).t().repeat(
            batch_size * self.num_anchors, 1, 1).view(y.shape).type(FloatTensor)

        # 生成先驗框的寬高
        anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
        anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
        anchor_w = anchor_w.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(w.shape)
        anchor_h = anchor_h.repeat(batch_size, 1).repeat(1, 1, input_height * input_width).view(h.shape)

        # 計算調整後的先驗框中心與寬高
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + grid_x
        pred_boxes[..., 1] = y.data + grid_y
        pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
        pred_boxes[..., 3] = torch.exp(h.data) * anchor_h

        # 用於將輸出調整爲相對於416x416的大小
        _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
        output = torch.cat((pred_boxes.view(batch_size, -1, 4) * _scale,
                            conf.view(batch_size, -1, 1), pred_cls.view(batch_size, -1, self.num_classes)), -1)
        return output.data

4、在原圖上進行繪製

通過第三步,我們可以獲得預測框在原圖上的位置,而且這些預測框都是經過篩選的。這些篩選後的框可以直接繪製在圖片上,就可以獲得結果了。

二、訓練部分

1、計算loss所需參數

在計算loss的時候,實際上是pred和target之間的對比:
pred就是網絡的預測結果。
target就是網絡的真實框情況。

2、pred是什麼

對於yolo3的模型來說,網絡最後輸出的內容就是三個特徵層每個網格點對應的預測框及其種類,即三個特徵層分別對應着圖片被分爲不同size的網格後,每個網格點上三個先驗框對應的位置、置信度及其種類。

輸出層的shape分別爲(13,13,75),(26,26,75),(52,52,75),最後一個維度爲75是因爲是基於voc數據集的,它的類爲20種,yolo3只有針對每一個特徵層存在3個先驗框,所以最後維度爲3x25;
如果使用的是coco訓練集,類則爲80種,最後的維度應該爲255 = 3x85
,三個特徵層的shape爲(13,13,255),(26,26,255),(52,52,255)

現在的y_pre還是沒有解碼的,解碼了之後纔是真實圖像上的情況。

3、target是什麼。

target就是一個真實圖像中,真實框的情況。
第一個維度是batch_size,第二個維度是每一張圖片裏面真實框的數量,第三個維度內部是真實框的信息,包括位置以及種類。

4、loss的計算過程

拿到pred和target後,不可以簡單的減一下作爲對比,需要進行如下步驟。

  1. 判斷真實框在圖片中的位置,判斷其屬於哪一個網格點去檢測。
  2. 判斷真實框和哪個先驗框重合程度最高。
  3. 計算該網格點應該有怎麼樣的預測結果才能獲得真實框
  4. 對所有真實框進行如上處理。
  5. 獲得網絡應該有的預測結果,將其與實際的預測結果對比。
from random import shuffle
import numpy as np
import torch
import torch.nn as nn
import math
import torch.nn.functional as F
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
from PIL import Image
from utils.utils import bbox_iou

def clip_by_tensor(t,t_min,t_max):
    t=t.float()
 
    result = (t >= t_min).float() * t + (t < t_min).float() * t_min
    result = (result <= t_max).float() * result + (result > t_max).float() * t_max
    return result

def MSELoss(pred,target):
    return (pred-target)**2

def BCELoss(pred,target):
    epsilon = 1e-7
    pred = clip_by_tensor(pred, epsilon, 1.0 - epsilon)
    output = -target * torch.log(pred) - (1.0 - target) * torch.log(1.0 - pred)
    return output

class YOLOLoss(nn.Module):
    def __init__(self, anchors, num_classes, img_size):
        super(YOLOLoss, self).__init__()
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.bbox_attrs = 5 + num_classes
        self.img_size = img_size

        self.ignore_threshold = 0.5
        self.lambda_xy = 1.0
        self.lambda_wh = 1.0
        self.lambda_conf = 1.0
        self.lambda_cls = 1.0

    def forward(self, input, targets=None):
        # 一共多少張圖片
        bs = input.size(0)
        # 特徵層的高
        in_h = input.size(2)
        # 特徵層的寬
        in_w = input.size(3)
        # 計算步長
        stride_h = self.img_size[1] / in_h
        stride_w = self.img_size[0] / in_w
        # 把先驗框的尺寸調整成特徵層大小的形式
        scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in self.anchors]
        # reshape
        prediction = input.view(bs, int(self.num_anchors/3),
                                self.bbox_attrs, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
        
        # 對prediction預測進行調整
        x = torch.sigmoid(prediction[..., 0])  # Center x
        y = torch.sigmoid(prediction[..., 1])  # Center y
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height
        conf = torch.sigmoid(prediction[..., 4])  # Conf
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

        # 找到哪些先驗框內部包含物體
        mask, noobj_mask, tx, ty, tw, th, tconf, tcls, box_loss_scale_x, box_loss_scale_y =\
                                                                            self.get_target(targets, scaled_anchors,
                                                                                            in_w, in_h,
                                                                                            self.ignore_threshold)

        noobj_mask = self.get_ignore(prediction, targets, scaled_anchors, in_w, in_h, noobj_mask)

        box_loss_scale_x = (2-box_loss_scale_x).cuda()
        box_loss_scale_y = (2-box_loss_scale_y).cuda()
        box_loss_scale = box_loss_scale_x*box_loss_scale_y
        mask, noobj_mask = mask.cuda(), noobj_mask.cuda()
        tx, ty, tw, th = tx.cuda(), ty.cuda(), tw.cuda(), th.cuda()
        tconf, tcls = tconf.cuda(), tcls.cuda()
        #  losses.
        loss_x = torch.sum(BCELoss(x, tx) / bs * box_loss_scale * mask)
        loss_y = torch.sum(BCELoss(y, ty) / bs * box_loss_scale * mask)
        loss_w = torch.sum(MSELoss(w, tw) / bs * 0.5 * box_loss_scale * mask)
        loss_h = torch.sum(MSELoss(h, th) / bs * 0.5 * box_loss_scale * mask)

        loss_conf = torch.sum(BCELoss(conf, mask) * mask / bs) + \
                    torch.sum(BCELoss(conf, mask) * noobj_mask / bs)
                    
        loss_cls = torch.sum(BCELoss(pred_cls[mask == 1], tcls[mask == 1])/bs)

        loss = loss_x * self.lambda_xy + loss_y * self.lambda_xy + \
                loss_w * self.lambda_wh + loss_h * self.lambda_wh + \
                loss_conf * self.lambda_conf + loss_cls * self.lambda_cls
        # print(loss, loss_x.item() + loss_y.item(), loss_w.item() + loss_h.item(), 
        #         loss_conf.item(), loss_cls.item(), \
        #         torch.sum(mask),torch.sum(noobj_mask))
        return loss, loss_x.item(), loss_y.item(), loss_w.item(), \
                loss_h.item(), loss_conf.item(), loss_cls.item()

    def get_target(self, target, anchors, in_w, in_h, ignore_threshold):
        # 計算一共有多少張圖片
        bs = len(target)
        # 獲得先驗框
        anchor_index = [[0,1,2],[3,4,5],[6,7,8]][[13,26,52].index(in_w)]
        subtract_index = [0,3,6][[13,26,52].index(in_w)]
        # 創建全是0或者全是1的陣列
        mask = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        noobj_mask = torch.ones(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)

        tx = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        ty = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        tw = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        th = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        tconf = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        tcls = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, self.num_classes, requires_grad=False)

        box_loss_scale_x = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        box_loss_scale_y = torch.zeros(bs, int(self.num_anchors/3), in_h, in_w, requires_grad=False)
        for b in range(bs):
            for t in range(target[b].shape[0]):
                # 計算出在特徵層上的點位
                gx = target[b][t, 0] * in_w
                gy = target[b][t, 1] * in_h
                gw = target[b][t, 2] * in_w
                gh = target[b][t, 3] * in_h

                # 計算出屬於哪個網格
                gi = int(gx)
                gj = int(gy)

                # 計算真實框的位置
                gt_box = torch.FloatTensor(np.array([0, 0, gw, gh])).unsqueeze(0)
                
                # 計算出所有先驗框的位置
                anchor_shapes = torch.FloatTensor(np.concatenate((np.zeros((self.num_anchors, 2)),
                                                                  np.array(anchors)), 1))
                # 計算重合程度
                anch_ious = bbox_iou(gt_box, anchor_shapes)
               
                # Find the best matching anchor box
                best_n = np.argmax(anch_ious)
                if best_n not in anchor_index:
                    continue
                # Masks
                if (gj < in_h) and (gi < in_w):
                    best_n = best_n - subtract_index
                    # 判定哪些先驗框內部真實的存在物體
                    noobj_mask[b, best_n, gj, gi] = 0
                    mask[b, best_n, gj, gi] = 1
                    # 計算先驗框中心調整參數
                    tx[b, best_n, gj, gi] = gx - gi
                    ty[b, best_n, gj, gi] = gy - gj
                    # 計算先驗框寬高調整參數
                    tw[b, best_n, gj, gi] = math.log(gw / anchors[best_n+subtract_index][0])
                    th[b, best_n, gj, gi] = math.log(gh / anchors[best_n+subtract_index][1])
                    # 用於獲得xywh的比例
                    box_loss_scale_x[b, best_n, gj, gi] = target[b][t, 2]
                    box_loss_scale_y[b, best_n, gj, gi] = target[b][t, 3]
                    # 物體置信度
                    tconf[b, best_n, gj, gi] = 1
                    # 種類
                    tcls[b, best_n, gj, gi, int(target[b][t, 4])] = 1
                else:
                    print('Step {0} out of bound'.format(b))
                    print('gj: {0}, height: {1} | gi: {2}, width: {3}'.format(gj, in_h, gi, in_w))
                    continue

        return mask, noobj_mask, tx, ty, tw, th, tconf, tcls, box_loss_scale_x, box_loss_scale_y

    def get_ignore(self,prediction,target,scaled_anchors,in_w, in_h,noobj_mask):
        bs = len(target)
        anchor_index = [[0,1,2],[3,4,5],[6,7,8]][[13,26,52].index(in_w)]
        scaled_anchors = np.array(scaled_anchors)[anchor_index]
        # print(scaled_anchors)
        # 先驗框的中心位置的調整參數
        x_all = torch.sigmoid(prediction[..., 0])  
        y_all = torch.sigmoid(prediction[..., 1])
        # 先驗框的寬高調整參數
        w_all = prediction[..., 2]  # Width
        h_all = prediction[..., 3]  # Height
        for i in range(bs):
            x = x_all[i]
            y = y_all[i]
            w = w_all[i]
            h = h_all[i]

            FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
            LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor

            # 生成網格,先驗框中心,網格左上角
            grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_w, 1).repeat(
                int(self.num_anchors/3), 1, 1).view(x.shape).type(FloatTensor)
            grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_h, 1).t().repeat(
                int(self.num_anchors/3), 1, 1).view(y.shape).type(FloatTensor)

            # 生成先驗框的寬高
            anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
            anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
            
            anchor_w = anchor_w.repeat(1, 1, in_h * in_w).view(w.shape)
            anchor_h = anchor_h.repeat(1, 1, in_h * in_w).view(h.shape)
            
            # 計算調整後的先驗框中心與寬高
            pred_boxes = torch.FloatTensor(prediction[0][..., :4].shape)
            pred_boxes[..., 0] = x.data + grid_x
            pred_boxes[..., 1] = y.data + grid_y
            pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
            pred_boxes[..., 3] = torch.exp(h.data) * anchor_h

            pred_boxes = pred_boxes.view(-1, 4)

            for t in range(target[i].shape[0]):
                gx = target[i][t, 0] * in_w
                gy = target[i][t, 1] * in_h
                gw = target[i][t, 2] * in_w
                gh = target[i][t, 3] * in_h
                gt_box = torch.FloatTensor(np.array([gx, gy, gw, gh])).unsqueeze(0)

                anch_ious = bbox_iou(gt_box, pred_boxes, x1y1x2y2=False)
                anch_ious = anch_ious.view(x.size())
                noobj_mask[i][anch_ious>self.ignore_threshold] = 0
                # print(torch.max(anch_ious))
        return noobj_mask

訓練自己的yolo3模型

yolo3整體的文件夾構架如下:
在這裏插入圖片描述
本文使用VOC格式進行訓練。
訓練前將標籤文件放在VOCdevkit文件夾下的VOC2007文件夾下的Annotation中。
在這裏插入圖片描述
訓練前將圖片文件放在VOCdevkit文件夾下的VOC2007文件夾下的JPEGImages中。
在這裏插入圖片描述
在訓練前利用voc2yolo3.py文件生成對應的txt。
在這裏插入圖片描述
再運行根目錄下的voc_annotation.py,運行前需要將classes改成你自己的classes。

classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

在這裏插入圖片描述
就會生成對應的2007_train.txt,每一行對應其圖片位置及其真實框的位置。
在這裏插入圖片描述
在訓練前需要修改model_data裏面的voc_classes.txt文件,需要將classes改成你自己的classes。同時還需要修改utils/config.py文件,修改內部的Num_Classes變成所分的種類的數量。
在這裏插入圖片描述
運行train.py即可開始訓練。
在這裏插入圖片描述

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章