【目標檢測】Faster R-CNN算法實現
一、前言
繼2014年的R-CNN、2015年的Fast R-CNN后，2016年目標檢測領域再次迎來Ross Girshick大佬的神作Faster R-CNN，一舉解決了目標檢測的實時性問題。相較于Fast R-CNN而言，Faster R-CNN主要改進措施包括：
- 區域建議網絡（RPN）：Faster R-CNN摒棄了選擇性搜索，提出區域提議網絡（Region Proposal Network）用于生成候選目標區域。RPN網絡通過滑窗的方式提取候選區域特征，并將提取的特征映射到一個低維向量，然后將其輸入到邊界框回歸層和邊界框分類層，獲取候選目標區域的位置偏移和分類輸出。
- 共享卷積特征：Faster R-CNN中RPN和檢測網絡共享輸入圖像的卷積特征，實現了端到端的聯合訓練，使得模型能夠更好地調整卷積特征以適應特定的檢測任務。
- 先驗框（Anchors）：Faster R-CNN中首次提出先驗框的概念，通過使用多尺度先驗框，RPN能夠生成不同大小和長寬比的候選區域，提高了模型對于不同尺度的目標的檢測能力。
上述改進措施使得Faster R-CNN在速度和準確性上都優于Fast R-CNN，它不僅具有更高的檢測精度，而且在處理多尺度和小目標問題時也更加有效。
同Fast RCNN實現一樣（見 http://www.diao-diao.com/Haitangr/p/17709548.html），本文將基于Pytorch框架，實現Faster RCNN算法，完成對17flowes數據集的花朵目標檢測任務。
二、Faster RCNN算法結構
Faster RCNN論文原文中算法整體結構如下：
如圖，Faster R-CNN算法流程主要包括四個部分，分別是卷積層（Conv Layers）、區域建議網絡（RPN）、感興趣區域池化（RoI Pool）和檢測網絡（Classifier）。各部分功能如下：
- 卷積層：卷積層是輸入圖像的特征提取器，作用是提取輸入圖像的全圖特征，用于RPN推薦區域生成和RoI區域池化。卷積層可以采用多種網絡結構來實現，比如Vgg/ResNet，本文代碼實現部分采用的是Vgg16部分結構；
- RPN：RPN作用是生成候選框。RPN在特征圖上滑動，對特征圖上每個位置生成一定數量的先驗框anchors，再結特征圖分類和回歸結果對先驗框位置進行校正，從中選取符合要求的候選區域proposals用于區域池化和目標檢測；
- 區域池化：區域池化的作用是根據proposals區域位置在特征圖上進行特征截取，并將截取的特征縮放到固定大小以便于檢測網絡進行分類和回歸；
- 檢測網絡：檢測網絡是最終對proposals進行分類和回歸的全連接網絡，用于對proposals進行更為精確的分類和位置回歸。
三、Faster RCNN算法實現
下圖為以VGG16為backbone的Faster RCNN模型結構。對于任意大小輸入圖像，在輸入網絡前會先統一縮放至MxN，然后經過VGG16進行提取高維特征圖。特征圖進入RPN網絡后有兩個分支流向，一是分類分支計算目標置信度，二是回歸分支計算位置偏移量，綜合二者結果獲得目標區域proposals。ROI Pooling層利用獲取的proposals從特征圖上提取區域特征，送入后續分類和回歸網絡進行最終預測。
1. Conv Layers
Conv Layers一般采用Vgg/ResNet結構作為backbone，目的是進行深度特征提取，將特征用于后續任務。本工程中使用的骨干網絡包含extractor和linear兩部分。其中extractor為Vgg16結構，包含4次2*2池化，其主要作用是對輸入圖像進行全圖特征提取。linear為全連接層結構，主要作用是對roi池化特征進行初步處理，便于最終檢測網絡分類和回歸。backbone具體實現如下：
def backbone(pretrained=False): """ 定義主干特征提取網絡和最終推薦區域特征處理線性網絡 :param pretrained: 是否加載預訓練參數 :return: 特征提取器和后續線性全連接層 """ net = torchvision.models.vgg16(pretrained=pretrained) extractor = net.features[:-1] linear = nn.Sequential( nn.Linear(in_features=512 * 7 * 7, out_features=4096, bias=True), nn.ReLU() ) if not pretrained: extractor.apply(lambda x: nn.init.kaiming_normal_(x.weight) if isinstance(x, nn.Conv2d) else None) linear.apply(lambda x: nn.init.kaiming_normal_(x.weight) if isinstance(x, nn.Linear) else None) return extractor, linear
2. RPN
RPN網絡是Faster RCNN的主要改進點，利用RPN網絡直接生成檢測框，極大的提高了檢測速度。RPN網絡示意圖如下
網絡主要分為兩個分支，即分類和回歸，分別計算目標概率和偏移值。實際代碼實現時主要分如下三步：
- 對特征圖像上每個位置生成固定數量一定尺寸的先驗框anchors，anchors會覆蓋圖像所有位置，并將進入后續網絡進行分類和回歸；
- 通過RPN中的分類和回歸分支判斷每個anchor是否包含目標并計算每個anchor的坐標偏移值；
- 在ProposalCreater中利用坐標偏移offsets對anchors進行位置校正，再通過尺寸、目標得分、非極大值抑制對anchors進行過濾，獲取最終的rois區域。
2.1 先驗框生成
先驗框anchors就是目標檢測中常用的一定尺寸的矩形框，每行有四個元素[x1, y1, x2, y2]，用來表示左上和右下角點坐標。在anchors生成過程中需要三組參數來控制anchors的大小和形狀，分別為base_size，wh_ratios和anchor_spatial_scales。其中base_size表示預設的基礎空間尺寸，默認為獲取feature_map過程中的空間縮放比例，表示特征圖上一個像素映射回原圖時的大小。wh_ratios表示錨框寬高比值，用來控制先驗框的形狀。anchor_spatial_scales表示錨框與基礎尺寸的比值，用來控制先驗框的大小。以base_size=16，wh_ratios=(0.5, 1.0, 2.0)，anchor_spatial_scales=(8, 16, 32)為例，算法會以原點為中心，按空間尺寸比分別生成邊長為16*8=128、16*16=256、16*32=512的正方形，然后在保持面積基本不變的情況下調整各邊長，得到寬高比分別為0.5、1.0、2.0的矩形先驗框。這樣三種空間尺寸和三種寬高比，能夠得到9個基礎先驗框。
基礎先驗框的具體代碼實現如下：
def generate_anchor_base(base_size=16, wh_ratios=(0.5, 1.0, 2.0), anchor_spatial_scales=(8, 16, 32)) -> np.ndarray: """ 生成基礎先驗框->(x1, y1, x2, y2) :param base_size: 預設基礎空間尺寸, 即特征圖與原圖的縮放比例, 表示映射回原圖時的空間大小 :param wh_ratios: 錨框寬高比w/h取值 :param anchor_spatial_scales: 待生成正方形框錨與預設的最基本錨框空間尺度上的縮放比例 :return: 生成的錨框 """ # 默認錨框左上角為(0, 0)時計算預設錨框的中心點坐標 cx = cy = (base_size - 1) / 2.0 # 根據錨框寬高比取值個數M1和錨框空間尺度縮放取值數量M2, 生成M2組面積基本相同但寬高比例不同的基礎錨框, 共N個（N=M1*M2） # 假設wh_ratios=(0.5, 1.0, 2.0)三種取值, anchor_scales=(8, 16, 32)三種取值, 那么生成的基礎錨框有9種可能取值 num_anchor_base = len(wh_ratios) * len(anchor_spatial_scales) # 生成[N, 4]維的基礎錨框 anchor_base = np.zeros((num_anchor_base, 4), dtype=np.float32) # 根據錨框面積計算錨框寬和高 # 錨框面積s=w*h, 而wh_ration=w/h, 則s=h*h*wh_ratio, 在已知面積和寬高比時: h=sqrt(s/wh_ratio) # 同樣可得s=w*w/wh_ratio, 在已知面積和寬高比時: w=sqrt(s*wh_ratio) # 計算不同寬高比、不同面積大小的錨框 for i in range(len(wh_ratios)): # 遍歷空間尺度縮放比例 for j in range(len(anchor_spatial_scales)): # 預設框面積為s1=base_size^2, 經空間尺度縮放后為s2=(base_size*anchor_spatial_scale)^2 # 將s2帶入上述錨框寬和高計算過程可求w和h值 h = base_size * anchor_spatial_scales[j] / np.sqrt(wh_ratios[i]) w = base_size * anchor_spatial_scales[j] * np.sqrt(wh_ratios[i]) idx = i * len(anchor_spatial_scales) + j anchor_base[idx, 0] = cx - (w - 1) / 2.0 anchor_base[idx, 1] = cy - (h - 1) / 2.0 anchor_base[idx, 2] = cx + (w - 1) / 2.0 anchor_base[idx, 3] = cy + (h - 1) / 2.0 return anchor_base
運行此代碼，可以生成9個基礎先驗框，這些框中的值都是映射回原圖后相對于（0，0）的坐標值
先驗框是基于特征圖上每個位置來生成的，特征圖有多上個點，就需要生成多少個先驗框。在上述9個基準先驗框的基礎上，可以通過坐標偏移來獲取全圖的先驗框，具體實現方法如下：
def generate_shifted_anchors(anchor_base: np.ndarray, fstep: int, fw: int, fh: int) -> np.ndarray: """ 根據基礎先驗框, 在特征圖上逐像素生成輸入圖像對應的先驗框 :param anchor_base: 預生成的基礎先驗框->[num_anchor_base, 4], 列坐標對應[x1, y1, x2, y2] :param fstep: 每個特征點映射回原圖后在原圖上的步進, 也就是空間縮放比例 :param fw: 特征圖像寬度 :param fh: 特征圖像高度 :return: 生成的全圖先驗框->[num_anchor_base * fw * fh, 4], 列坐標對應[x1, y1, x2, y2] """ # 特征圖上每個點都會生成anchor, 第一個特征點對應原圖上的先驗框就是anchor_base, 由于特征圖與原圖間空間縮放, 相鄰兩特征點對應的anchor在原圖上的步進為圖像空間縮放尺度大小fstep # 因此由anchor_base和anchor間步進, 可以計算出輸入圖像的所有anchor # 計算出原圖每一行上anchor與anchor_base的位置偏移值取值 shift_x = np.arange(0, fw * fstep, fstep) # 計算出原圖每一列上anchor與anchor_base的位置偏移值取值 shift_y = np.arange(0, fh * fstep, fstep) # 用兩方向偏移值生成網格, x和y方向偏移值維度均為[fh, fw], shift_x/shift_y相同位置的兩個值表示當前anchor相對于anchor_base的坐標偏移值 shift_x, shift_y = np.meshgrid(shift_x, shift_y) # 將shift_x/shift_y展開成一維并按列拼接在一起, 分別對應anchor的x1/y1/x2/y2的坐標偏移值, 構成偏移值矩陣 shift = np.stack([shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel()], axis=1) # anchor_base維度為[num_anchor_base, 4], 偏移值矩陣shift維度為[fw * fh, 4] # 將anchor_base每一行與shift中每一行元素相加即得到位置偏移后的所有anchor坐標, 但二者維度不同需進行擴展 num_anchor_base = anchor_base.shape[0] num_points = shift.shape[0] # 將兩者均擴展為[num_points, num_anchors, 4] anchor_base_extend = np.repeat(a=anchor_base[np.newaxis, :, :], repeats=num_points, axis=0) shift_extend = np.repeat(a=shift[:, np.newaxis, :], repeats=num_anchor_base, axis=1) # 獲取最終anchors坐標, 并展開成二維向量, 維度為[num_anchors * num_points, 4] = [num_anchors * fw * fh, 4] anchors = anchor_base_extend + shift_extend anchors = np.reshape(a=anchors, newshape=(-1, 4)) return anchors
2.2 前景判斷與偏移值計算
特征圖進入RPN網絡后會先經過一個3*3卷積，進一步集中特征信息，之后分別流向分類分支和回歸分支。
在分類分支中，輸入特征（預設512通道）經過1*1卷積后，變成18通道特征，維度為[num，18，fh，fw]。其中num為輸入圖像數量，fh和fw為特征圖高和寬。之所以設計為18通道，是因為每個特征點生成了9個形狀大小不一的anchors，每個anchors有前景和背景兩種可能，所以每個特征點對應的anchors有18種結果。這樣18通道和18種結果相互對應，經過后面softmax處理后，能夠根據置信度挑選出目標區域。之后通過reshape操作，經過[num，num_anchor_base * 2，fh，fw]→[num，fh，fw，num_anchor_base * 2]→[num，fh * fw * num_anchor_base，2]的維度變換，能夠使分類結果更便于后續使用。
同理，在回歸分支中，輸入特征經過1*1卷積后，變成36通道特征，維度為[num，36，fh，fw]。其中36對應每個位置9個先驗框，每個先驗框有4個坐標值。為了后續使用方便，本文工程實現過程中也對結果進行了[num，num_anchor_base * 4，fh，fw]→[num，fh，fw，num_anchor_base * 4]→[num，fh * fw * num_anchor_base，4]的維度變換。
前景判斷和偏移值計算部分代碼如下：
# rpn網絡的3*3卷積, 目的是對輸入feature map進行卷積, 進一步集中特征信息 self.conv = nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=(3, 3), stride=(1, 1), padding=1) # rpn網絡分類分支, 逐像素對特征圖上anchor進行分類(每個anchor對應前景/背景兩類, 每個類別為一個通道, 總通道數為num_anchor_base * 2) self.classifier = nn.Conv2d(in_channels=mid_channels, out_channels=num_anchor_base * 2, kernel_size=(1, 1), stride=(1, 1), padding=0) # rpn網絡回歸分支, 逐像素對特征圖上anchor進行坐標回歸(每個框對應四個坐標值, 每個坐標為一個通道, 總通道數為num_anchor_base * 4) self.regressor = nn.Conv2d(in_channels=mid_channels, out_channels=num_anchor_base * 4, kernel_size=(1, 1), stride=(1, 1), padding=0) self.relu = nn.ReLU() self.softmax = nn.Softmax(dim=-1)
num, chans, fh, fw = x.size() # 先將輸入圖像特征經過3*3網絡, 進一步特征處理用于后續分類和回歸 x = self.conv(x) x = self.relu(x) # 將特征圖送入分類網絡, 計算特征圖上每個像素的分類結果, 對應原輸入圖像所有先驗框的分類結果, 維度為[num, num_anchor_base * 2, fh, fw] out_classifier = self.classifier(x) # 維度轉換[num, num_anchor_base * 2, fh, fw]->[num, fh, fw, num_anchor_base * 2]->[num, fh * fw * num_anchor_base, 2] out_classifier = out_classifier.permute(0, 2, 3, 1).contiguous().view(num, -1, 2) # 將特征圖送入回歸網絡, 計算特征圖上每個像素的回歸結果, 對應原輸入圖像所有先驗框的回歸結果, 維度為[num, num_anchor_base * 4, fh, fw] out_offsets = self.regressor(x) # 維度轉換[num, num_anchor_base * 4, fh, fw]->[num, fh, fw, num_anchor_base * 4]->[num, fh * fw * num_anchor_base, 4] out_offsets = out_offsets.permute(0, 2, 3, 1).contiguous().view(num, -1, 4) # 將分類器輸出轉換為得分 out_scores = self.softmax(out_classifier) # out_scores[:, :, 1]表示存在目標的概率 out_scores = out_scores[:, :, 1].contiguous().view(num, -1)
2.3 目標區域過濾
按照上述流程能夠在一幅圖像上提取到大量先驗框。假設圖像縮放值[640, 800]后輸入VGG16網絡，經過特征提取后，其尺寸變為[40，50]，那么能夠得到40*50*9=18000個先驗框，但是這些先驗框存在三個問題，一是部分框坐標超出原圖像范圍，二是很多框內不存在待檢測目標，三是這些框存在大量重疊，因此需要對其進一步過濾篩選。具體過濾流程依次為以下幾個步驟：
- 利用2.2中計算出的偏移值offsets對2.1中計算出的先驗框anchors位置進行修正，計算位置校正后的anchors；
- 剔除超出圖像邊界范圍的anchors；
- 限定最小尺寸min_size，剔除尺寸過小的anchors；
- 根據2.2中前景得分對anchors進行降序排列，保留得分最高前N個anchors；
- 對最后保留的N個anchors進行nms非極大值抑制，保留的rois作為過濾后的目標區域
本文中目標區域過濾采用ProposalCreator來實現，代碼如下：
class ProposalCreator: # 對每一幅圖像, 利用偏移值對所有先驗框進行位置矯正得到目標建議框, 再通過尺寸限制/得分限制/nms方法對目標建議框進行過濾, 獲得推薦區域, 即每幅圖像的roi區域 def __init__(self, nms_thresh=0.7, num_samples_train=(12000, 2000), num_samples_test=(6000, 300), min_size=16, train_flag=False): """ 初始化推薦區域生成器, 為每幅圖像生成滿足尺寸要求、得分要求、nms要求的規定數量推薦框 :param nms_thresh: 非極大值抑制閾值 :param num_samples_train: 訓練過程非極大值抑制前后待保留的樣本數 :param num_samples_test: 測試過程非極大值抑制步驟前后待保留的樣本數 :param min_size: 邊界框最小寬高限制 :param train_flag: 模型訓練還是測試 """ self.train_flag = train_flag self.nms_thresh = nms_thresh self.num_samples_train = num_samples_train self.num_samples_test = num_samples_test self.min_size = min_size @staticmethod def calc_bboxes_from_offsets(offsets: Tensor, anchors: Tensor, eps=1e-5) -> Tensor: """ 由圖像特征計算的偏移值offsets對rpn產生的先驗框位置進行修正 :param offsets: 偏移值矩陣->[n, 4], 列對應[x1, y1, x2, y2]的偏移值 :param anchors: 先驗框矩陣->[n, 4], 列坐標對應[x1, y1, x2, y2] :param eps: 極小值, 防止乘以0或者負數 :return: 目標坐標矩陣->[n, 4], 對應[x1, y1, x2, y2] """ eps = torch.tensor(eps).type_as(offsets) targets = torch.zeros_like(offsets, dtype=torch.float32) # 計算目標真值框中心點坐標及長寬 anchors_h = anchors[:, 3] - anchors[:, 1] + 1 anchors_w = anchors[:, 2] - anchors[:, 0] + 1 anchors_cx = 0.5 * (anchors[:, 2] + anchors[:, 0]) anchors_cy = 0.5 * (anchors[:, 1] + anchors[:, 3]) anchors_w = torch.maximum(anchors_w, eps) anchors_h = torch.maximum(anchors_h, eps) # 將偏移值疊加到真值上計算anchors的中心點和寬高 targets_w = anchors_w * torch.exp(offsets[:, 2]) targets_h = anchors_h * torch.exp(offsets[:, 3]) targets_cx = anchors_cx + offsets[:, 0] * anchors_w targets_cy = anchors_cy + offsets[:, 1] * anchors_h targets[:, 0] = targets_cx - 0.5 * (targets_w - 1) targets[:, 1] = targets_cy - 0.5 * (targets_h - 1) targets[:, 2] = targets_cx + 0.5 * (targets_w - 1) targets[:, 3] = targets_cy + 0.5 * (targets_h - 1) return targets def __call__(self, offsets: Tensor, anchors: Tensor, scores: Tensor, im_size: tuple, scale: float = 1.0) -> Tensor: """ 利用回歸器偏移值/全圖先驗框/分類器得分生成滿足條件的推薦區域 :param offsets: 偏移值->[fw * fh * num_anchor_base, 4], 列坐標對應[x1, y1, x2, y2] :param anchors: 全圖先驗框->[fw * fh * num_anchor_base, 4], 列坐標對應[x1, y1, x2, y2] :param scores: 目標得分->[fw * fh * num_anchor_base] :param im_size: 原始輸入圖像大小 (im_height, im_width) :param scale: scale和min_size一起控制先驗框最小尺寸 :return: 經過偏移值矯正及過濾后保留的目標建議框, 維度[num_samples_after_nms, 4]列坐標對應[x1, y1, x2, y2] """ # 設置nms過程前后需保留的樣本數量, 注意訓練和測試過程保留的樣本數量不一致 if self.train_flag: num_samples_before_nms, num_samples_after_nms = self.num_samples_train else: num_samples_before_nms, num_samples_after_nms = self.num_samples_test im_height, im_width = im_size # 利用回歸器計算的偏移值對全圖先驗框位置進行矯正, 獲取矯正后的目標先驗框坐標 targets = self.calc_bboxes_from_offsets(offsets=offsets, anchors=anchors) # 對目標目標先驗框坐標進行限制, 防止坐標落在圖像外 # 保證0 <= [x1, x2] <= cols - 1 targets[:, [0, 2]] = torch.clip(targets[:, [0, 2]], min=0, max=im_width - 1) # 0 <= [y1, y2] <= rows - 1 targets[:, [1, 3]] = torch.clip(targets[:, [1, 3]], min=0, max=im_height - 1) # 利用min_size和scale控制先驗框尺寸下限, 移除尺寸太小的目標先驗框 min_size = self.min_size * scale # 計算目標先驗框寬高 targets_w = targets[:, 2] - targets[:, 0] + 1 targets_h = targets[:, 3] - targets[:, 1] + 1 # 根據寬高判斷框是否有效, 挑選出有效邊界框和得分 is_valid = (targets_w >= min_size) & (targets_h >= min_size) targets = targets[is_valid] scores = scores[is_valid] # 利用區域目標得分對目標先驗框數量進行限制 # 對目標得分進行降序排列, 獲取降序索引 descend_order = torch.argsort(input=scores, descending=True) # 在nms之前, 選取固定數量得分稍高的目標先驗框 if num_samples_before_nms > 0: descend_order = descend_order[:num_samples_before_nms] targets = targets[descend_order] scores = scores[descend_order] # 利用非極大值抑制限制邊界框數量 keep = nms(boxes=targets, scores=scores, iou_threshold=self.nms_thresh) # 如果數量不足則隨機抽取, 用于填補不足 if len(keep) < num_samples_after_nms: random_indexes = np.random.choice(a=range(len(keep)), size=(num_samples_after_nms - len(keep)), replace=True) keep = torch.concat([keep, keep[random_indexes]]) # 在nms后, 截取固定數量邊界框, 即最終生成的roi區域 keep = keep[:num_samples_after_nms] targets = targets[keep] return targets
上述代碼利用num_samples_train和num_samples_test來控制訓練與預測過程中的rois數量。訓練過程中每幅圖像生成2000個roi區域，預測過程中每幅圖像生成300個roi區域。
2.4 RPN實現
本文中RPN網絡整體結構由RegionProposalNet來定義實現，具體代碼如下：
class RegionProposalNet(nn.Module): # rpn推薦區域生成網絡, 獲取由圖像特征前向計算得到的回歸器偏移值/分類器結果, 以及經過偏移值矯正后的roi區域/roi對應數據索引/全圖先驗框 def __init__(self, in_channels=512, mid_channels=512, feature_stride=16, wh_ratios=(0.5, 1.0, 2.0), anchor_spatial_scales=(8, 16, 32), train_flag=False): super(RegionProposalNet, self).__init__() # 特征點步長, 即特征提取網絡的空間尺度縮放比例, 工程中vgg16網絡使用了4層2*2的MaxPool, 故取值為16 self.feature_stride = feature_stride # 推薦框生成器: 對全圖先驗框進行位置矯正和過濾, 提取位置矯正后的固定數量的目標推薦框 self.create_proposals = ProposalCreator(train_flag=train_flag) # 根據寬高比和空間尺度比例, 生成固定數量的基礎先驗框 self.anchor_base = self.generate_anchor_base(wh_ratios=wh_ratios, anchor_spatial_scales=anchor_spatial_scales) num_anchor_base = len(wh_ratios) * len(anchor_spatial_scales) # rpn網絡的3*3卷積, 目的是對輸入feature map進行卷積, 進一步集中特征信息 self.conv = nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=(3, 3), stride=(1, 1), padding=1) # rpn網絡分類分支, 逐像素對特征圖上anchor進行分類(每個anchor對應前景/背景兩類, 每個類別為一個通道, 總通道數為num_anchor_base * 2) self.classifier = nn.Conv2d(in_channels=mid_channels, out_channels=num_anchor_base * 2, kernel_size=(1, 1), stride=(1, 1), padding=0) # rpn網絡回歸分支, 逐像素對特征圖上anchor進行坐標回歸(每個框對應四個坐標值, 每個坐標為一個通道, 總通道數為num_anchor_base * 4) self.regressor = nn.Conv2d(in_channels=mid_channels, out_channels=num_anchor_base * 4, kernel_size=(1, 1), stride=(1, 1), padding=0) self.relu = nn.ReLU() self.softmax = nn.Softmax(dim=-1) # 權重初始化 nn.init.kaiming_normal_(self.conv.weight) nn.init.constant_(self.conv.bias, 0.0) nn.init.kaiming_normal_(self.classifier.weight) nn.init.constant_(self.classifier.bias, 0.0) nn.init.kaiming_normal_(self.regressor.weight) nn.init.constant_(self.regressor.bias, 0.0) @staticmethod def generate_anchor_base(base_size=16, wh_ratios=(0.5, 1.0, 2.0), anchor_spatial_scales=(8, 16, 32)) -> np.ndarray: """ 生成基礎先驗框->(x1, y1, x2, y2) :param base_size: 預設的最基本的正方形錨框邊長 :param wh_ratios: 錨框寬高比w/h取值 :param anchor_spatial_scales: 待生成正方形框錨與預設的最基本錨框空間尺度上的縮放比例 :return: 生成的錨框 """ # 默認錨框左上角為(0, 0)時計算預設錨框的中心點坐標 cx = cy = (base_size - 1) / 2.0 # 根據錨框寬高比取值個數M1和錨框空間尺度縮放取值數量M2, 生成M2組面積基本相同但寬高比例不同的基礎錨框, 共N個（N=M1*M2） # 假設wh_ratios=(0.5, 1.0, 2.0)三種取值, anchor_scales=(8, 16, 32)三種取值, 那么生成的基礎錨框有9種可能取值 num_anchor_base = len(wh_ratios) * len(anchor_spatial_scales) # 生成[N, 4]維的基礎錨框 anchor_base = np.zeros((num_anchor_base, 4), dtype=np.float32) # 根據錨框面積計算錨框寬和高 # 錨框面積s=w*h, 而wh_ration=w/h, 則s=h*h*wh_ratio, 在已知面積和寬高比時: h=sqrt(s/wh_ratio) # 同樣可得s=w*w/wh_ratio, 在已知面積和寬高比時: w=sqrt(s*wh_ratio) # 計算不同寬高比、不同面積大小的錨框 for i in range(len(wh_ratios)): # 遍歷空間尺度縮放比例 for j in range(len(anchor_spatial_scales)): # 預設框面積為s1=base_size^2, 經空間尺度縮放后為s2=(base_size*anchor_spatial_scale)^2 # 將s2帶入上述錨框寬和高計算過程可求w和h值 h = base_size * anchor_spatial_scales[j] / np.sqrt(wh_ratios[i]) w = base_size * anchor_spatial_scales[j] * np.sqrt(wh_ratios[i]) idx = i * len(anchor_spatial_scales) + j anchor_base[idx, 0] = cx - (w - 1) / 2.0 anchor_base[idx, 1] = cy - (h - 1) / 2.0 anchor_base[idx, 2] = cx + (w - 1) / 2.0 anchor_base[idx, 3] = cy + (h - 1) / 2.0 return anchor_base @staticmethod def generate_shifted_anchors(anchor_base: np.ndarray, fstep: int, fw: int, fh: int) -> np.ndarray: """ 根據基礎先驗框, 在特征圖上逐像素生成輸入圖像對應的先驗框 :param anchor_base: 預生成的基礎先驗框->[num_anchor_base, 4], 列坐標對應[x1, y1, x2, y2] :param fstep: 每個特征點映射回原圖后在原圖上的步進, 也就是空間縮放比例 :param fw: 特征圖像寬度 :param fh: 特征圖像高度 :return: 生成的全圖先驗框->[num_anchor_base * fw * fh, 4], 列坐標對應[x1, y1, x2, y2] """ # 特征圖上每個點都會生成anchor, 第一個特征點對應原圖上的先驗框就是anchor_base, 由于特征圖與原圖間空間縮放, 相鄰兩特征點對應的anchor在原圖上的步進為圖像空間縮放尺度大小fstep # 因此由anchor_base和anchor間步進, 可以計算出輸入圖像的所有anchor # 計算出原圖每一行上anchor與anchor_base的位置偏移值取值 shift_x = np.arange(0, fw * fstep, fstep) # 計算出原圖每一列上anchor與anchor_base的位置偏移值取值 shift_y = np.arange(0, fh * fstep, fstep) # 用兩方向偏移值生成網格, x和y方向偏移值維度均為[fh, fw], shift_x/shift_y相同位置的兩個值表示當前anchor相對于anchor_base的坐標偏移值 shift_x, shift_y = np.meshgrid(shift_x, shift_y) # 將shift_x/shift_y展開成一維并按列拼接在一起, 分別對應anchor的x1/y1/x2/y2的坐標偏移值, 構成偏移值矩陣 shift = np.stack([shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel()], axis=1) # anchor_base維度為[num_anchor_base, 4], 偏移值矩陣shift維度為[fw * fh, 4] # 將anchor_base每一行與shift中每一行元素相加即得到位置偏移后的所有anchor坐標, 但二者維度不同需進行擴展 num_anchor_base = anchor_base.shape[0] num_points = shift.shape[0] # 將兩者均擴展為[num_points, num_anchors, 4] anchor_base_extend = np.repeat(a=anchor_base[np.newaxis, :, :], repeats=num_points, axis=0) shift_extend = np.repeat(a=shift[:, np.newaxis, :], repeats=num_anchor_base, axis=1) # 獲取最終anchors坐標, 并展開成二維向量, 維度為[num_anchors * num_points, 4] = [num_anchors * fw * fh, 4] anchors = anchor_base_extend + shift_extend anchors = np.reshape(a=anchors, newshape=(-1, 4)) return anchors def forward(self, x: Tensor, im_size: tuple, scale: float = 1.0) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """ 前向處理, 獲取rpn網絡的回歸器輸出/分類器輸出/矯正的roi區域/roi數據索引/全圖先驗框 :param x: 由原圖提取的輸入特征圖feature_map->[num, 512, fh, fw] :param im_size: 原始輸入圖像尺寸->[im_width, im_height] :param scale: 與min_size一起用于控制最小先驗框尺寸 :return: list->[回歸器輸出, 分類器輸出, 建議框, 建議框對應的數據索引, 全圖先驗框] """ num, chans, fh, fw = x.size() # 先將輸入圖像特征經過3*3網絡, 進一步特征處理用于后續分類和回歸 x = self.conv(x) x = self.relu(x) # 將特征圖送入分類網絡, 計算特征圖上每個像素的分類結果, 對應原輸入圖像所有先驗框的分類結果, 維度為[num, num_anchor_base * 2, fh, fw] out_classifier = self.classifier(x) # 維度轉換[num, num_anchor_base * 2, fh, fw]->[num, fh, fw, num_anchor_base * 2]->[num, fh * fw * num_anchor_base, 2] out_classifier = out_classifier.permute(0, 2, 3, 1).contiguous().view(num, -1, 2) # 將特征圖送入回歸網絡, 計算特征圖上每個像素的回歸結果, 對應原輸入圖像所有先驗框的回歸結果, 維度為[num, num_anchor_base * 4, fh, fw] out_offsets = self.regressor(x) # 維度轉換[num, num_anchor_base * 4, fh, fw]->[num, fh, fw, num_anchor_base * 4]->[num, fh * fw * num_anchor_base, 4] out_offsets = out_offsets.permute(0, 2, 3, 1).contiguous().view(num, -1, 4) # 將分類器輸出轉換為得分 out_scores = self.softmax(out_classifier) # out_scores[:, :, 1]表示存在目標的概率 out_scores = out_scores[:, :, 1].contiguous().view(num, -1) # 生成全圖先驗框, 每個特征點都會生成num_anchor_base個先驗框, 故全圖生成的先驗框維度為[fh * fw * num_anchor_base, 4] anchors = self.generate_shifted_anchors(anchor_base=self.anchor_base, fstep=self.feature_stride, fw=fw, fh=fh) # 將anchors轉換到和out_offsets的數據類型(float32)和設備類型(cuda/cpu)一致 anchors = torch.tensor(data=anchors).type_as(out_offsets) # 獲取batch數據的roi區域和對應的索引 rois, rois_idx = [], [] # 遍歷batch中每個數據 for i in range(num): # 按照數據索引獲取當前數據對應的偏移值和得分, 生成固定數量的推薦區域, 維度為[fixed_num, 4] proposals = self.create_proposals(offsets=out_offsets[i], anchors=anchors, scores=out_scores[i], im_size=im_size, scale=scale) # 創建和推薦區域proposals數量相同的batch索引, 維度為[fixed_num] batch_idx = torch.tensor([i] * len(proposals)) # 將推薦區域和索引進行維度擴展后放入列表中, 擴展后二者維度分別為[1 ,fixed_num, 4]和[[1 ,fixed_num]] rois.append(proposals.unsqueeze(0)) rois_idx.append(batch_idx.unsqueeze(0)) # 將rois列表中所有張量沿0維拼接在一起. 原rois列表長度為num, 其中每個張量均為[1, fixed_num, 4], 拼接后rois張量維度為[num, fixed_num, 4] rois = torch.cat(rois, dim=0).type_as(x) # 將rois索引拼接在一起, 拼接后維度為[num, fixed_num] rois_idx = torch.cat(rois_idx, dim=0).to(x.device) return out_offsets, out_classifier, rois, rois_idx, anchors
3. ROI Pooling與Classifier
RPN網絡獲得的rois目標區域，其大小并不是固定的，對應到特征圖上也就是不同尺寸。但后續區域分類和位置回歸需要輸入固定大小特征數據進入全連接網絡，所以需要對每個roi區域提取的特征進行尺寸縮放，使其變換為同一大小。此外，經過ROI Pooling后，獲得的區域池化特征pool_features又有兩個流向，一是進行邊界框位置預測，二是進行目標類別判斷。
注意，Classifier與RPN都包含分類和回歸兩個分支，雖然這兩部分結構共享卷積特征feature map，有助于減少計算量并提高檢測速度，但二者仍具有明顯的區別：
- 結構不同：RPN中分類和回歸分支是利用卷積，用對用通道表示分類和回歸結果，Classifier使用全連接結果來表示分類和回歸結果；
- 任務不同：RPN中分類分支負責對生成的anchor進行二分類，判斷其是前景還是背景，任務相對簡單，目的是篩選出可能包含目標物體的proposals?；貧w分支則負責對這些anchor的位置進行微調，使其更接近于真實的目標邊界框，關注的是對候選區域位置的初步修正。Classifier中的分類分支需要對RPN網絡篩選出的候選區域進行多類別分類，確定每個區域具體屬于哪個物體類別，任務更復雜。同時，Classifier中的回歸分支則負責對候選區域的位置進行進一步的精細調整，以獲得更準確的目標定位，關注的是對目標物體位置的精確預測。
本文中采用ROIHead來實現ROI Pooling與Classifier，代碼如下：
class ROIHead(nn.Module): def __init__(self, num_classes: int, pool_size: int, linear: nn.Module, spatial_scale: float = 1.0): """ 將ROI區域送入模型獲得分類器輸出和回歸器輸出 :param num_classes: 樣本類別數 :param pool_size: roi池化目標尺寸 :param linear: 線性模型 :param spatial_scale: roi池化所使用的空間比例, 默認1.0, 若待處理的roi坐標為原圖坐標, 則此處需要設置spatial_scale=目標特征圖大小/原圖大小 """ super(ROIHead, self).__init__() self.linear = linear # 對roi_pool結果進行回歸預測 self.regressor = nn.Linear(4096, num_classes * 4) self.classifier = nn.Linear(4096, num_classes) nn.init.kaiming_normal_(self.regressor.weight) nn.init.kaiming_normal_(self.classifier.weight) # 后續采用的roi坐標為特征圖上坐標, 因此spatial_scale直接設置為1.0即可 # 注意roi_pool要求roi坐標滿足格式[x1, y1, x2, y2] self.roi_pool = RoIPool(output_size=(pool_size, pool_size), spatial_scale=spatial_scale) def forward(self, x: Tensor, rois: Tensor, rois_idx: Tensor, im_size: Tensor) -> Tuple[Tensor, Tensor]: """ 根據推薦框對特征圖進行roi池化, 并將結果送入分類器和回歸器, 得到相應結果 :param x: 輸入batch數據對應的全圖圖像特征, 維度為[num, 512, fh, fw] :param rois: 輸入batch數據對應的rois區域, 維度為[num, num_samples, 4], 順序為[y1, x1, y2, x2] :param rois_idx: 輸入batch數據對應的rois區域索引, 維度為[num, num_samples] :param im_size: 原始輸入圖像尺寸, 維度為[im_height, im_width] :return: 分類模型和回歸模型結果 """ num, chans, fh, fw = x.size() im_height, im_width = im_size # 將一個batch內數據的推薦區域展開堆疊在一起, 維度變為[num * num_samples, 4] rois = torch.flatten(input=rois, start_dim=0, end_dim=1) # 將一個batch內數據的索引展開堆疊在一起, 維度變為[num * num_samples] rois_idx = torch.flatten(input=rois_idx, start_dim=0, end_dim=1) # 計算原圖roi區域映射到特征圖后對應的邊界框位置, 維度為[num * num_samples, 4] rois_on_features = torch.zeros_like(rois) # 計算[x1, x2]映射后的坐標 rois_on_features[:, [0, 2]] = rois[:, [0, 2]] * (fw / im_width) # 計算[y1, y2]映射后的坐標 rois_on_features[:, [1, 3]] = rois[:, [1, 3]] * (fh / im_height) # 將特征圖上roi區域和對應的數據索引在列方向進行拼接, 得到[num * num_samples, 5]維張量, 用于后續roi_pool, 列對應[idx, x1, y1, x2, y2] fidx_and_rois = torch.cat([rois_idx.unsqueeze(1), rois_on_features], dim=1) # 根據數據idx和推薦框roi對輸入圖像特征圖進行截取 # 注意由于rois_on_features中roi坐標已經縮放到了特征圖大小, 所以RoIPool池化時的spatial_scale需要設置為1.0 # 注意此處roi_pool需要num * num_samples, 5]維, 根據第0列的idx取x中截取相應的特征進行池化 pool_features = self.roi_pool(x, fidx_and_rois) # 上面獲取的池化特征維度為[num * num_samples, chans, pool_size, pool_size], 將其展開為[num * num_samples, chans * pool_size * pool_size]以便送入分類和回歸網絡 pool_features = pool_features.view(pool_features.size(0), -1) # 利用線性層進行特征進一步濃縮 linear_features = self.linear(pool_features) # 將樣本特征送入回歸器, 得到各樣本輸出, 維度為[num * num_samples, 4 * num_classes] rois_out_regressor = self.regressor(linear_features) # 將樣本特征送入回歸器, 得到各樣本輸出, 維度為[num * num_samples, num_classes] rois_out_classifier = self.classifier(linear_features) # 維度變換, 獲得當前batch中每個數據的所有回歸結果, 維度為[num, num_samples, 4 * num_classes] rois_out_regressor = rois_out_regressor.view(num, -1, rois_out_regressor.size(1)) # 維度變換, 獲得當前batch中每個數據的所有分類結果, 維度為[num, num_samples, num_classes] rois_out_classifier = rois_out_classifier.view(num, -1, rois_out_classifier.size(1)) return rois_out_regressor, rois_out_classifier
4. Faster RCNN代碼
由上可知，本文中Faster RCNN包含backbone、RPN、ROI Head三大主體結構，其具體代碼實現如下：
class FasterRCNN(nn.Module): def __init__(self, num_classes, train_flag=False, feature_stride=16, anchor_spatial_scales=(8, 16, 32), wh_ratios=(0.5, 1.0, 2.0), pretrained=False): """ 初始化Faster R-CNN :param num_classes: 最終分類類別數, 包含背景0和目標類別數 :param train_flag: 是否為訓練過程 :param feature_stride: 特征步進, 實際就是特征提取器的空間縮放比例, 工程使用移除最后一個池化層的vgg16, 特征空間縮放比例為16 :param anchor_spatial_scales: 待生成先驗框與基本先驗框的邊長比值 :param wh_ratios: 待生成先驗框的寬高比 :param pretrained: 特征提取器是否加載預訓練參數 """ super(FasterRCNN, self).__init__() self.feature_stride = feature_stride self.extractor, linear = backbone(pretrained=pretrained) self.rpn = RegionProposalNet(in_channels=512, mid_channels=512, feature_stride=feature_stride, wh_ratios=wh_ratios, anchor_spatial_scales=anchor_spatial_scales, train_flag=train_flag) self.head = ROIHead(num_classes=num_classes, pool_size=POOL_SIZE, spatial_scale=1, linear=linear) def forward(self, x, scale: float = 1.0, mode: str = "forward"): """ Faster R-CNN前向過程 :param x: 輸入 :param scale: rpn結構中用于控制最小先驗框尺寸 :param mode: 處理流程控制字符串 :return: """ if mode == "forward": im_size = x.size()[-2:] # 提取輸入圖像特征 im_features = self.extractor(x) # 獲取建議框 _, _, rois, rois_idx, _ = self.rpn(im_features, im_size, scale) # 根據圖像特征和建議框計算偏移值回歸結果和區域分類結果 rois_out_regressor, rois_out_classifier = self.head(im_features, rois, rois_idx, im_size) return rois_out_regressor, rois_out_classifier, rois, rois_idx elif mode == "extractor": # 提取圖像特征 im_features = self.extractor(x) return im_features elif mode == "rpn": im_features, im_size = x # 獲取建議框 out_offsets, out_classes, rois, rois_idx, anchors = self.rpn(im_features, im_size, scale) return out_offsets, out_classes, rois, rois_idx, anchors elif mode == "head": im_features, rois, rois_idx, im_size = x # 獲取分類和回歸結果 rois_out_regressor, rois_out_classifier = self.head(im_features, rois, rois_idx, im_size) return rois_out_regressor, rois_out_classifier else: raise TypeError("Invalid parameter of mode, which must be in ['forward', 'extractor', 'rpn', 'head']")
至此，Faster RCNN網絡結構實現基本完成，之后將開始模型訓練和預測。
四、模型訓練
本文Faster RCNN模型訓練按照以下順序進行：
- 計算輸入batch圖像數據特征im_features；
- 將im_features輸入RPN網絡，獲取感興趣區域rois并計算RPN網絡分類和回歸損失；
- 對im_features中感興趣區域進行ROI Pooling，將結果送入Classifier計算rois對應的分類和回歸損失；
- 計算整體損失并反向傳播，優化網絡參數。
1. 獲取圖像特征
輸入圖像經backbone的extractor（VGG16）進行多次卷積池化，可以得到輸入批數據特征im_features，具體實現如下：
if step < train_batch_num: # 設置為訓練模式 network.train() # 獲取輸入圖像全圖特征, 維度為[num, 512, im_height/16, im_width/16] im_features = network.forward(x=ims, mode="extractor")
2. RPN訓練
2.1 RPN輸出
將圖像輸入RPN網絡后，可以獲取rois以及RPN分類和回歸輸出，具體實現如下：
# 利用rpn網絡獲取回歸器輸出/分類器輸出/batch數據對應建議框/建議框對應數據索引/全圖先驗框 rpn_offsets, rpn_classifier, rois, rois_idx, anchors = network.forward(x=[im_features, (im_height, im_width)], mode="rpn") np_anchors = anchors.cpu().numpy()
此處除了rois、rpn_offsets、rpn_classifier外，還保留了rois_idx和anchors結果。其中anchors是先驗框，后續獲取RPN網絡分類和回歸真值時需要使用，rois_idx存儲的是每個roi區域對應的batch數據中的圖片索引。
2.2 RPN真值
我們已經獲得RPN分類和回歸分支輸出，訓練網絡需要損失函數，也就還需要得到對應真值。對于RPN分類分支，其真值對應的是每個anchor的標簽，但實際我們只有每個圖象上真實邊界框gt_boxes對應的標簽，我們肯定不能用gt_boxes標簽來定義anchor標簽，否則若一幅圖上某個gt_box標簽為1，那么所有anchor標簽也為1，整幅圖上所有位置都標記為1了，這顯然是不合理的。應該如何來定義每個anchor對應的標簽呢？具體方案如下：
- 剔除不在圖像邊界范圍內的anchors；
- 計算anchors與gt_boxes的交并比iou；
- 尋找每個gt_box下iou最大時的anchor索引max_iou_idx_anchor，每個anchor對應iou最大時的gt_box索引max_iou_idx_gt以及當前anchor對應的iou最大值max_iou_values_anchor；
- 如果max_iou_values_anchor小于預設的負樣本iou閾值negative_iou_thresh，表示當前anchor與所有gt_box重疊面積小，應置為背景（標簽0）；
- 對于max_iou_idx_anchor，表示與當前gt_box重疊最大的anchor索引，即最接近真實目標位置的anchor，應置為前景（標簽1）；
- 同樣，如果max_iou_values_anchor大于預設的正樣本iou閾值positive_iou_thresh，表示當前anchor與所有gt_box重疊面積大，應置為前景（標簽1）；
- 對于其他情況，統一將標簽設置為-1，表示既不是正樣本也不是負樣本，不用于RPN網絡訓練；
- 獲得的正負樣本數量可能較大也可能數量不足，故還需要對其正負樣本數量進行一定限制。
通過上面的方案獲取了分類網絡中anchors對應的標簽labels，還需要獲取回歸網絡anchors對應的偏移值真值offsets。上面我們已經求得每個anchor對應iou最大時的gt_box索引為max_iou_idx_gt，那么對應的這個gt_box就理應作為當前anchor的位置真值，據此能夠獲得偏移值作為RPN回歸網絡真值。
這個過程用AnchorCreator來定義實現，具體代碼如下：
class AnchorCreator: # 生成先驗框對應的標簽及與真值框間的真實偏移值 def __init__(self, num_samples=256, positive_iou_thresh=0.7, negative_iou_thresh=0.3, positive_rate=0.5): """ 初始化anchor生成器 :param num_samples: 每幀圖片上用于后續分類和回歸任務的有效推薦區域總數 :param positive_iou_thresh: 正樣本的IoU判定閾值 :param negative_iou_thresh: 負樣本的判定閾值 :param positive_rate: 正樣本所占樣本總數的比例 """ self.num_samples = num_samples self.positive_iou_thresh = positive_iou_thresh self.negative_iou_thresh = negative_iou_thresh self.positive_rate = positive_rate @staticmethod def is_inside_anchors(anchors: Union[np.ndarray, Tensor], width: int, height: int) -> Union[np.ndarray, Tensor]: """ 獲取圖像內部的推薦框 :param anchors: 生成的所有推薦框->[x1, y1, x2, y2] :param width: 輸入圖像寬度 :param height: 輸入圖像高度 :return: 未超出圖像邊界的推薦框 """ is_inside = (anchors[:, 0] >= 0) & (anchors[:, 1] >= 0) & (anchors[:, 2] <= width - 1) & (anchors[:, 3] <= height - 1) return is_inside @staticmethod def calc_IoU(anchors: np.ndarray, gt_boxes: np.ndarray, method=1) -> np.ndarray: """ 計算推薦區域與真值的IoU :param anchors: 推薦區域邊界框, [m, 4]維數組, 四列分別對應左上和右下兩個點坐標[x1, y1, x2, y2] :param gt_boxes: 當前圖像中所有真值邊界框, [n, 4]維數組, 四列分別對應左上和右下兩點坐標[x1, y1, x2, y2] :param method: iou計算方法 :return: iou, [m, n]維數組, 記錄每個推薦區域與每個真值框的IoU結果 """ # 先判斷維度是否符合要求 assert anchors.ndim == gt_boxes.ndim == 2, "anchors and ground truth bbox must be 2D array." assert anchors.shape[1] == gt_boxes.shape[1] == 4, "anchors and ground truth bbox must contain 4 values for 2 points." num_anchors, num_gts = anchors.shape[0], gt_boxes.shape[0] # 方法1: 利用for循環遍歷求解交并比 if method == 0: iou = np.zeros((num_anchors, num_gts)) # anchor有m個, gt_box有n個, 遍歷求出每個gt_box對應的iou結果即可 for idx in range(num_gts): gt_box = gt_boxes[idx] box_area = (anchors[:, 2] - anchors[:, 0]) * (anchors[:, 3] - anchors[:, 1]) gt_area = (gt_box[2] - gt_box[0]) * (gt_box[3] - gt_box[1]) inter_w = np.minimum(anchors[:, 2], gt_box[2]) - np.maximum(anchors[:, 0], gt_box[0]) inter_h = np.minimum(anchors[:, 3], gt_box[3]) - np.maximum(anchors[:, 1], gt_box[1]) inter = np.maximum(inter_w, 0) * np.maximum(inter_h, 0) union = box_area + gt_area - inter iou[:, idx] = inter / union # 方法2: 利用repeat對矩陣進行升維, 從而利用對應位置計算交并比 elif method == 1: # anchors維度為[m, 4], gt_boxes維度為[n, 4], 對二者通過repeat的方式都升維到[m, n, 4] anchors = np.repeat(anchors[:, np.newaxis, :], num_gts, axis=1) gt_boxes = np.repeat(gt_boxes[np.newaxis, :, :], num_anchors, axis=0) # 利用對應位置求解框面積 anchors_area = (anchors[:, :, 2] - anchors[:, :, 0]) * (anchors[:, :, 3] - anchors[:, :, 1]) gt_boxes_area = (gt_boxes[:, :, 2] - gt_boxes[:, :, 0]) * (gt_boxes[:, :, 3] - gt_boxes[:, :, 1]) # 求交集區域的寬和高 inter_w = np.minimum(anchors[:, :, 2], gt_boxes[:, :, 2]) - np.maximum(anchors[:, :, 0], gt_boxes[:, :, 0]) inter_h = np.minimum(anchors[:, :, 3], gt_boxes[:, :, 3]) - np.maximum(anchors[:, :, 1], gt_boxes[:, :, 1]) # 求交并比 inter = np.maximum(inter_w, 0) * np.maximum(inter_h, 0) union = anchors_area + gt_boxes_area - inter iou = inter / union # 方法3: 利用np函數的廣播機制求結果而避免使用循環 else: # 計算anchors和gt_boxes左上角點的最大值, 包括兩x1坐標最大值和y1坐標最大值 # 注意anchors[:, None, :2]會增加一個新維度, 維度為[m, 1, 2], gt_boxes[:, :2]維度為[n, 2], maximum計算最大值時會將二者都擴展到[m, n, 2] max_left_top = np.maximum(anchors[:, None, :2], gt_boxes[:, :2]) # 計算anchors和gt_boxes右下角點的最小值, 包括兩x2坐標最大值和y2坐標最大值, 同上也用到了廣播機制 min_right_bottom = np.minimum(anchors[:, None, 2:], gt_boxes[:, 2:]) # 求交集面積和并集面積 # min_right_bottom - max_left_top維度為[m, n, 2], 后兩列代表交集區域的寬和高 # 用product進行兩列元素乘積求交集面積, 用(max_left_top < min_right_bottom).all(axis=2)判斷寬和高是否大于0, 結果維度為[m, n] inter = np.product(min_right_bottom - max_left_top, axis=2) * (max_left_top < min_right_bottom).all(axis=2) # 用product進行兩列元素乘積求每個anchor的面積, 結果維度維[m] anchors_area = np.product(anchors[:, 2:] - anchors[:, :2], axis=1) # 用product進行兩列元素乘積求每個gt_box的面積, 結果維度維[n] gt_boxes_area = np.product(gt_boxes[:, 2:] - gt_boxes[:, :2], axis=1) # anchors_area[:, None]維度維[m, 1], gt_boxes_area維度維[n], 二者先廣播到[m, n]維度, 再和同緯度inter做減法計算, 結果維度維[m, n] union = anchors_area[:, None] + gt_boxes_area - inter iou = inter / union return iou @staticmethod def calc_max_iou_info(iou: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ 利用iou結果計算出最大iou及其對應位置 :param iou: [m, n]維矩陣, 其中m為anchors數量, n為gt_boxes數量 :return: 每一列最大iou出現的行編號, 每一行最大iou出現的列編號, 每一行的最大iou結果 """ # 按列求每一列的iou最大值出現的行數, 即記錄與每個gt_box的iou最大的anchor的行編號, 維度和gt_box個數相同, 為n(每個gt_box對應一個anchor與之iou最大) max_iou_idx_anchor = np.argmax(iou, axis=0) # 按行求每一行的iou最大值出現的列數, 即記錄與每個anchor的iou最大的gt_box的列編號, 維度和anchor個數相同, 為m(每個anchor對應一個gt_box與之iou最大) max_iou_idx_gt = np.argmax(iou, axis=1) # 求每個anchor與所有gt_box的最大iou值 max_iou_values_anchor = np.max(iou, axis=1) return max_iou_idx_anchor, max_iou_idx_gt, max_iou_values_anchor def create_anchor_labels(self, anchors: np.ndarray, gt_boxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ 計算IoU結果并根據結果為每個推薦區域生成標簽 :param anchors: 生成的有效推薦區域, 列坐標對應[x1, y1, x2, y2] :param gt_boxes: 真值框, 列坐標對應[x1, y1, x2, y2] :return: 每個推薦區域的最大iou對應的真值框編號, 推薦區域對應的標簽 """ # 計算iou結果 iou = self.calc_IoU(anchors=anchors, gt_boxes=gt_boxes) # 計算行/列方向最大iou對應位置和值 max_iou_idx_anchor, max_iou_idx_gt, max_iou_values_anchor = self.calc_max_iou_info(iou=iou) # 先將所有label置為-1, -1表示不進行處理, 既不是正樣本也不是負樣本, 再根據iou判定正樣本為1, 背景為0 labels = -1 * np.ones(anchors.shape[0], dtype="int") # max_iou_values_anchor為每一行最大的iou結果, 其值低于負樣本閾值, 表明該行對應的anchor與所有gt_boxes的iou結果均小于閾值, 設置為負樣本 labels[max_iou_values_anchor < self.negative_iou_thresh] = 0 # max_iou_idx_anchor為每一列iou最大值出現的行編號, 表明對應行anchor與某個gt_box的iou最大, iou最大肯定是設置為正樣本 labels[max_iou_idx_anchor] = 1 # max_iou_values_anchor為每一行最大的iou結果, 其值大于正樣本閾值, 表明該行對應的anchor與至少一個gt_box的iou結果大于閾值, 設置為正樣本 labels[max_iou_values_anchor >= self.positive_iou_thresh] = 1 # 對正負樣本數量進行限制 # 計算目標正樣本數量 num_positive = int(self.num_samples * self.positive_rate) # 記錄正樣本行編號 idx_positive = np.where(labels == 1)[0] if len(idx_positive) > num_positive: size_to_rest = len(idx_positive) - num_positive # 從正樣本編號中隨機選取一定數量將標簽置為-1 idx_to_reset = np.random.choice(a=idx_positive, size=size_to_rest, replace=False) labels[idx_to_reset] = -1 # 計算現有負樣本數量 num_negative = self.num_samples - np.sum(labels == 1) # 記錄負樣本行編號 idx_negative = np.where(labels == 0)[0] if len(idx_negative) > num_negative: size_to_rest = len(idx_negative) - num_negative # 從負樣本編號中隨機選取一定數量將標簽置為-1 idx_to_reset = np.random.choice(a=idx_negative, size=size_to_rest, replace=False) labels[idx_to_reset] = -1 return max_iou_idx_gt, labels @staticmethod def calc_offsets_from_bboxes(anchors: np.ndarray, target_boxes: np.ndarray, eps: float = 1e-5) -> np.ndarray: """ 計算推薦區域與真值間的位置偏移 :param anchors: 候選邊界框, 列坐標對應[x1, y1, x2, y2] :param target_boxes: 真值, 列坐標對應[x1, y1, x2, y2] :param eps: 極小值, 防止除以0或者負數 :return: 邊界框偏移值->[dx, dy, dw, dh] """ offsets = np.zeros_like(anchors, dtype="float32") # 計算anchor中心點坐標及長寬 anchors_h = anchors[:, 3] - anchors[:, 1] + 1 anchors_w = anchors[:, 2] - anchors[:, 0] + 1 anchors_cy = 0.5 * (anchors[:, 3] + anchors[:, 1]) anchors_cx = 0.5 * (anchors[:, 2] + anchors[:, 0]) # 計算目標真值框中心點坐標及長寬 targets_h = target_boxes[:, 3] - target_boxes[:, 1] + 1 targets_w = target_boxes[:, 2] - target_boxes[:, 0] + 1 targets_cy = 0.5 * (target_boxes[:, 3] + target_boxes[:, 1]) targets_cx = 0.5 * (target_boxes[:, 2] + target_boxes[:, 0]) # 限制anchor長寬防止小于0 anchors_w = np.maximum(anchors_w, eps) anchors_h = np.maximum(anchors_h, eps) # 計算偏移值 offsets[:, 0] = (targets_cx - anchors_cx) / anchors_w offsets[:, 1] = (targets_cy - anchors_cy) / anchors_h offsets[:, 2] = np.log(targets_w / anchors_w) offsets[:, 3] = np.log(targets_h / anchors_h) return offsets def __call__(self, im_width: int, im_height: int, anchors: np.ndarray, gt_boxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ 利用真值框和先驗框的iou結果為每個先驗框打標簽, 同時計算先驗框和真值框對應的偏移值 :param im_width: 輸入圖像寬度 :param im_height: 輸入圖像高度 :param anchors: 全圖先驗框, 列坐標對應[x1, y1, x2, y2] :param gt_boxes: 真值框, 列坐標對應[x1, y1, x2, y2] :return: 先驗框對應標簽和應該產生的偏移值 """ num_anchors = len(anchors) # 獲取有效的推薦區域, 其維度為[m], m <= num_anchors is_inside = self.is_inside_anchors(anchors=anchors, width=im_width, height=im_height) inside_anchors = anchors[is_inside] # 在有效先驗框基礎上, 獲取每個先驗框的最大iou對應的真值框編號和區域標簽 max_iou_idx_gt, inside_labels = self.create_anchor_labels(anchors=inside_anchors, gt_boxes=gt_boxes) # 每個anchor都存在n個真值框, 選擇最大iou對應的那個真值框作為每個anchor的目標框計算位置偏移 # gt_boxes維度為[n, 4], max_iou_idx_gt維度為[m], 從真值中挑選m次即得到與每個anchor的iou最大的真值框, 即所需目標框, 維度為[m, 4] target_boxes = gt_boxes[max_iou_idx_gt] inside_offsets = self.calc_offsets_from_bboxes(anchors=inside_anchors, target_boxes=target_boxes) # 上面的偏移值和labels都是在inside_anchors中求得, 現在將結果映射回全圖 # 將所有標簽先置為-1, 再將內部先驗框標簽映射回全圖 labels = -1 * np.ones(num_anchors) labels[is_inside] = inside_labels # 將所有偏移值先置為0, 再將內部先驗框偏移值映射回全圖 offsets = np.zeros_like(anchors) offsets[is_inside] = inside_offsets return labels, offsets
在訓練過程中通過遍歷的方式，計算每張圖片RPN中對應的真值，代碼如下：
anchor_creator = AnchorCreator() # 遍歷每一個數據的真值框/數據標簽/rpn網絡輸出 for i in range(num): # 獲取每張圖像的真值框/標簽/rpn生成的偏移值/rpn生成的分類輸出 cur_gt_boxes = gt_boxes[i] cur_labels = labels[i] cur_rpn_offsets = rpn_offsets[i] cur_rpn_classifier = rpn_classifier[i] cur_rois = rois[i] np_cur_gt_boxes = cur_gt_boxes.clone().detach().cpu().numpy() np_cur_rois = cur_rois.clone().detach().cpu().numpy() np_cur_labels = cur_labels.clone().detach().cpu().numpy() # 根據當前圖像真值框和先驗框, 獲取每個先驗框標簽以及圖像經過rpn網絡后應產生的偏移值作為真值 cur_gt_rpn_labels, cur_gt_rpn_offsets = anchor_creator(im_width=im_width, im_height=im_height, anchors=np_anchors, gt_boxes=np_cur_gt_boxes)
2.3 RPN損失
獲得RPN輸出和真值后，可以計算相應的損失值。RPN分為分類和回歸兩個任務，每個任務都需要進行計算。代碼實現如下：
def multitask_loss(out_offsets: Tensor, out_classifier: Tensor, gt_offsets: Tensor, gt_labels: Tensor, alpha: float = 1.0) -> Tuple[Tensor, Tensor, Tensor]: """ 計算多任務損失 :param out_offsets: 回歸模型邊界框結果 :param out_classifier: 分類模型邊界框結果 :param gt_offsets: 真實邊界框 :param gt_labels: 邊界框標簽 :param alpha: 權重系數 :return: 分類損失/正樣本回歸損失/總損失 """ # 分類損失計算式忽略標簽為-1的樣本 cls_loss_func = nn.CrossEntropyLoss(ignore_index=-1) reg_loss_func = nn.SmoothL1Loss() # 計算分類損失 loss_cls = cls_loss_func(out_classifier, gt_labels) # 選擇正樣本計算回歸損失 out_offsets_valid = out_offsets[gt_labels > 0] gt_offsets_valid = gt_offsets[gt_labels > 0] loss_reg = reg_loss_func(out_offsets_valid, gt_offsets_valid) # 總損失 loss = loss_cls + alpha * loss_reg return loss_cls, loss_reg, loss
在計算損失時，標簽為-1既不是正樣本也不是負樣本，需要忽略，所以計算回歸損失時設置ignore_index=-1，表明忽略標簽-1，計算回歸損失時，也采用設置gt_labels>0。
3. Classifier訓練
3.1 roi區域篩選
經過RPN網絡后，每張訓練圖片能生成2000個roi區域用于精確分類和定位，這些區域并不是全部用于ROI Pooling和最終的Classifier網絡。我們會對這些區域進一步篩選，并計算篩選后的結果與真值之間的偏移，作為最終Classifier網絡的真值用于Classifier網絡訓練。具體實施方法如下：
- 計算推薦區域rois與真值gt_boxes的交并比iou；
- 計算每個roi對應iou最大時的gt_box索引max_iou_idx_gt以及每個roi對應的iou最大值max_iou_values_anchor；
- 根據max_iou_idx_gt獲取rois對應的真值框位置roi_gt_boxes和對應的標簽roi_gt_labels；
- 通過隨機抽取的方式控制保留的rois數量，得到用于后續任務的樣本sample_rois；
- 計算保留的sample_rois與對應真值框sample_gt_boxes的位置偏移sample_offsets，以及對應的多分類標簽sample_labels，作為最終Classifier網絡的真值。
上述過程定義ProposalTargetCreator來實現，經過處理后，2000個roi區域最終會保留128個用于Classifier訓練。具體代碼如下：
class ProposalTargetCreator: def __init__(self, num_samples=128, positive_iou_thresh=0.5, negative_iou_thresh=(0.5, 0.0), positive_rate=0.5): """ 在roi區域中選擇一定數量的正負樣本區域, 計算坐標偏移和分類標簽, 用于后續分類和回歸網絡 :param num_samples: 待保留的正負樣本總數 :param positive_iou_thresh: 正樣本閾值 :param negative_iou_thresh: 負樣本閾值最大值和最小值 :param positive_rate: 正樣本比例 """ self.num_samples = num_samples self.positive_iou_thresh = positive_iou_thresh self.negative_iou_thresh = negative_iou_thresh self.positive_rate = positive_rate self.num_positive_per_image = int(num_samples * positive_rate) # 定義坐標偏移值歸一化系數, 用于正負樣本區域的offsets歸一化 self.offsets_normalize_params = np.array([[0, 0, 0, 0], [0.1, 0.1, 0.2, 0.2]], dtype="float32") def __call__(self, rois: np.ndarray, gt_boxes: np.ndarray, labels: np.ndarray): """ 根據推薦區域的iou結果選擇一定數量的推薦區域作為正負樣本, 并計算推薦區域與真值間的偏移值 :param rois: 推薦區域, 維度為[m, 4] :param gt_boxes: 真值框, 維度為[n, 4] :param labels: 圖像類別標簽, 維度為[l, 1], 注意此處取值為[1, num_target_classes], 默認背景為0 :return: 保留的正負樣本區域/區域偏移值/區域標簽 """ rois = np.concatenate((rois, gt_boxes), axis=0) # 計算iou結果 iou = AnchorCreator.calc_IoU(anchors=rois, gt_boxes=gt_boxes) # 根據iou最大獲取每個推薦框對應的真實框的idx和相應iou結果 _, max_iou_idx_gt, max_iou_values = AnchorCreator.calc_max_iou_info(iou=iou) # 獲取每個roi區域對應的真值框 roi_gt_boxes = gt_boxes[max_iou_idx_gt] # 獲取每個roi區域對應的真值框標簽, 取值從1開始, 如果取值從0開始, 由于存在背景, 真值框標簽需要額外加1 roi_gt_labels = labels[max_iou_idx_gt] # 選取roi區域中的正樣本序號, np.where()返回滿足條件的元組, 元組第一個元素為行索引, 第二個元素為列索引 positive_idx = np.where(max_iou_values >= self.positive_iou_thresh)[0] num_positive = min(self.num_positive_per_image, len(positive_idx)) if len(positive_idx) > 1: positive_idx = np.random.choice(a=positive_idx, size=num_positive, replace=False) # 選取roi區域中的負樣本序號 negative_idx = np.where((max_iou_values < self.negative_iou_thresh[0]) & (max_iou_values >= self.negative_iou_thresh[1]))[0] num_negative = min(self.num_samples - num_positive, len(negative_idx)) if len(negative_idx) > 1: negative_idx = np.random.choice(a=negative_idx, size=num_negative, replace=False) # 將正負樣本索引整合在一起, 獲得所有樣本索引 sample_idx = np.append(positive_idx, negative_idx) # 提取正負樣本對應的真值標簽, 此時無論正/負roi_gt_labels中都為對應iou最大的真值框標簽, 下一步就需要對負樣本標簽賦值為0 sample_labels = roi_gt_labels[sample_idx] # 對正負樣本中的負樣本標簽賦值為0 sample_labels[num_positive:] = 0 # 提取樣本對應的roi區域 sample_rois = rois[sample_idx] # 計算選取的樣本roi與真值的坐標偏移值 # 根據樣本索引, 獲取樣本對應的真值框 sample_gt_boxes = roi_gt_boxes[sample_idx] # 計算推薦區域樣本與真值的坐標偏移 sample_offsets = AnchorCreator.calc_offsets_from_bboxes(anchors=sample_rois, target_boxes=sample_gt_boxes) # 對坐標偏移進行歸一化 sample_offsets = (sample_offsets - self.offsets_normalize_params[0]) / self.offsets_normalize_params[1] return sample_rois, sample_offsets, sample_labels
訓練過程中，遍歷每張圖片獲取最終推薦區域后，需要將每個batch的數據結果整合在一起，其代碼如下：
# 在當前圖像生成的roi建議框中中, 抽取一定數量的正負樣本, 并計算出相應位置偏移, 用于后續回歸和區域分類 sample_rois, sample_offsets, sample_labels = proposal_creator(rois=np_cur_rois, gt_boxes=np_cur_gt_boxes, labels=np_cur_labels) # 將每個圖像生成的樣本信息存儲起來用于后續回歸和分類 # 抽取當前數據生成的推薦區域樣本放入list中, list長度為num, 每個元素維度為[num_samples, 4] samples_rois.append(torch.tensor(sample_rois).type_as(rpn_offsets)) # 將抽取的樣本索引放入list中, list長度為num, 每個元素維度為[num_samples] samples_indexes.append(torch.ones(len(sample_rois), device=rpn_offsets.device) * rois_idx[i][0]) # 將抽取的樣本偏移值放入list中, list長度為num, 每個元素維度為[num_samples, 4] samples_offsets.append(torch.tensor(sample_offsets).type_as(rpn_offsets)) # 將抽取的樣本分類標簽放入list中, list長度為num, 每個元素維度為[num_samples] samples_labels.append(torch.tensor(sample_labels, device=rpn_offsets.device).long())
3.2 Classifier損失
將圖像特征im_features、保留的推薦區域sample_rois經過ROIHead處理，可以得到每張圖片Classifier網絡的分類和回歸輸出。結合4.3.1中獲取的真值，即可計算該網絡的損失。代碼如下：
# 整合當前batch數據抽取的推薦區域信息 samples_rois = torch.stack(samples_rois, dim=0) samples_indexes = torch.stack(samples_indexes, dim=0) # 將圖像特征和推薦區域送入模型, 進行roi池化并獲得分類模型/回歸模型輸出 roi_out_offsets, roi_out_classifier = network.forward(x=[im_features, samples_rois, samples_indexes, (im_height, im_width)], mode="head") # 遍歷每幀圖像的roi信息 for i in range(num): cur_num_samples = roi_out_offsets.size(1) cur_roi_out_offsets = roi_out_offsets[i] cur_roi_out_classifier = roi_out_classifier[i] cur_roi_gt_offsets = samples_offsets[i] cur_roi_gt_labels = samples_labels[i] # 將當前數據的roi區域由[cur_num_samples, num_classes * 4] -> [cur_num_samples, num_classes, 4] cur_roi_out_offsets = cur_roi_out_offsets.view(cur_num_samples, -1, 4) # 根據roi對應的樣本標簽, 選擇與其類別對應真實框的offsets cur_roi_offsets = cur_roi_out_offsets[torch.arange(0, cur_num_samples), cur_roi_gt_labels] # 計算分類網絡和回歸網絡的損失值 cur_roi_cls_loss, cur_roi_reg_loss, _ = multitask_loss(out_offsets=cur_roi_offsets, out_classifier=cur_roi_out_classifier, gt_offsets=cur_roi_gt_offsets, gt_labels=cur_roi_gt_labels) roi_cls_loss += cur_roi_cls_loss roi_reg_loss += cur_roi_reg_loss
自此，Classifier網絡訓練結束。
4. 訓練代碼
綜合上述流程，得到最終的訓練部分代碼Train.py：
import os import torch import numpy as np import torch.nn as nn from torch import Tensor from typing import Tuple import matplotlib.pyplot as plt from Utils import GenDataSet from torch.optim.lr_scheduler import StepLR from Model import AnchorCreator, ProposalTargetCreator, FasterRCNN from Config import * def multitask_loss(out_offsets: Tensor, out_classifier: Tensor, gt_offsets: Tensor, gt_labels: Tensor, alpha: float = 1.0) -> Tuple[Tensor, Tensor, Tensor]: """ 計算多任務損失 :param out_offsets: 回歸模型邊界框結果 :param out_classifier: 分類模型邊界框結果 :param gt_offsets: 真實邊界框 :param gt_labels: 邊界框標簽 :param alpha: 權重系數 :return: 分類損失/正樣本回歸損失/總損失 """ # 分類損失計算式忽略標簽為-1的樣本 cls_loss_func = nn.CrossEntropyLoss(ignore_index=-1) reg_loss_func = nn.SmoothL1Loss() # 計算分類損失 loss_cls = cls_loss_func(out_classifier, gt_labels) # 選擇正樣本計算回歸損失 out_offsets_valid = out_offsets[gt_labels > 0] gt_offsets_valid = gt_offsets[gt_labels > 0] loss_reg = reg_loss_func(out_offsets_valid, gt_offsets_valid) # 總損失 loss = loss_cls + alpha * loss_reg return loss_cls, loss_reg, loss def train(data_set, network, num_epochs, optimizer, scheduler, device, train_rate: float = 0.8): """ 模型訓練 :param data_set: 訓練數據集 :param network: 網絡結構 :param num_epochs: 訓練輪次 :param optimizer: 優化器 :param scheduler: 學習率調度器 :param device: CPU/GPU :param train_rate: 訓練集比例 :return: None """ os.makedirs('./model', exist_ok=True) network = network.to(device) best_loss = np.inf print("=" * 8 + "開始訓練模型" + "=" * 8) # 計算訓練batch數量 batch_num = len(data_set) train_batch_num = round(batch_num * train_rate) # 記錄訓練過程中每一輪損失和準確率 train_loss_all, val_loss_all, train_acc_all, val_acc_all = [], [], [], [] anchor_creator = AnchorCreator() proposal_creator = ProposalTargetCreator() for epoch in range(num_epochs): # 記錄train/val總損失 num_train_loss = num_val_loss = train_loss = val_loss = 0.0 for step, batch_data in enumerate(data_set): # 讀取數據, 注意gt_boxes列坐標對應[x1, y1, x2, y2] ims, labels, gt_boxes = batch_data ims = ims.to(device) labels = labels.to(device) gt_boxes = gt_boxes.to(device) num, chans, im_height, im_width = ims.size() if step < train_batch_num: # 設置為訓練模式 network.train() # 獲取輸入圖像全圖特征, 維度為[num, 512, im_height/16, im_width/16] im_features = network.forward(x=ims, mode="extractor") # 利用rpn網絡獲取回歸器輸出/分類器輸出/batch數據對應建議框/建議框對應數據索引/全圖先驗框 rpn_offsets, rpn_classifier, rois, rois_idx, anchors = network.forward(x=[im_features, (im_height, im_width)], mode="rpn") np_anchors = anchors.cpu().numpy() # 記錄rpn區域推薦網絡的分類/回歸損失, 以及最終的roi區域分類和回歸損失 rpn_cls_loss, rpn_reg_loss, roi_cls_loss, roi_reg_loss = 0, 0, 0, 0 samples_rois, samples_indexes, samples_offsets, samples_labels = [], [], [], [] # 遍歷每一個數據的真值框/數據標簽/rpn網絡輸出 for i in range(num): # 獲取每張圖像的真值框/標簽/rpn生成的偏移值/rpn生成的分類輸出 cur_gt_boxes = gt_boxes[i] cur_labels = labels[i] cur_rpn_offsets = rpn_offsets[i] cur_rpn_classifier = rpn_classifier[i] cur_rois = rois[i] np_cur_gt_boxes = cur_gt_boxes.clone().detach().cpu().numpy() np_cur_rois = cur_rois.clone().detach().cpu().numpy() np_cur_labels = cur_labels.clone().detach().cpu().numpy() # 根據當前圖像真值框和先驗框, 獲取每個先驗框標簽以及圖像經過rpn網絡后應產生的偏移值作為真值 cur_gt_rpn_labels, cur_gt_rpn_offsets = anchor_creator(im_width=im_width, im_height=im_height, anchors=np_anchors, gt_boxes=np_cur_gt_boxes) # 轉換為張量后計算rpn網絡的回歸損失和分類損失 cur_gt_rpn_offsets = torch.tensor(cur_gt_rpn_offsets).type_as(rpn_offsets) cur_gt_rpn_labels = torch.tensor(cur_gt_rpn_labels).long().to(rpn_offsets.device) cur_rpn_cls_loss, cur_rpn_reg_loss, _ = multitask_loss(out_offsets=cur_rpn_offsets, out_classifier=cur_rpn_classifier, gt_offsets=cur_gt_rpn_offsets, gt_labels=cur_gt_rpn_labels) rpn_cls_loss += cur_rpn_cls_loss rpn_reg_loss += cur_rpn_reg_loss # 在當前圖像生成的roi建議框中中, 抽取一定數量的正負樣本, 并計算出相應位置偏移, 用于后續回歸和區域分類 sample_rois, sample_offsets, sample_labels = proposal_creator(rois=np_cur_rois, gt_boxes=np_cur_gt_boxes, labels=np_cur_labels) # 將每個圖像生成的樣本信息存儲起來用于后續回歸和分類 # 抽取當前數據生成的推薦區域樣本放入list中, list長度為num, 每個元素維度為[num_samples, 4] samples_rois.append(torch.tensor(sample_rois).type_as(rpn_offsets)) # 將抽取的樣本索引放入list中, list長度為num, 每個元素維度為[num_samples] samples_indexes.append(torch.ones(len(sample_rois), device=rpn_offsets.device) * rois_idx[i][0]) # 將抽取的樣本偏移值放入list中, list長度為num, 每個元素維度為[num_samples, 4] samples_offsets.append(torch.tensor(sample_offsets).type_as(rpn_offsets)) # 將抽取的樣本分類標簽放入list中, list長度為num, 每個元素維度為[num_samples] samples_labels.append(torch.tensor(sample_labels, device=rpn_offsets.device).long()) # 整合當前batch數據抽取的推薦區域信息 samples_rois = torch.stack(samples_rois, dim=0) samples_indexes = torch.stack(samples_indexes, dim=0) # 將圖像特征和推薦區域送入模型, 進行roi池化并獲得分類模型/回歸模型輸出 roi_out_offsets, roi_out_classifier = network.forward(x=[im_features, samples_rois, samples_indexes, (im_height, im_width)], mode="head") # 遍歷每幀圖像的roi信息 for i in range(num): cur_num_samples = roi_out_offsets.size(1) cur_roi_out_offsets = roi_out_offsets[i] cur_roi_out_classifier = roi_out_classifier[i] cur_roi_gt_offsets = samples_offsets[i] cur_roi_gt_labels = samples_labels[i] # 將當前數據的roi區域由[cur_num_samples, num_classes * 4] -> [cur_num_samples, num_classes, 4] cur_roi_out_offsets = cur_roi_out_offsets.view(cur_num_samples, -1, 4) # 根據roi對應的樣本標簽, 選擇與其類別對應真實框的offsets cur_roi_offsets = cur_roi_out_offsets[torch.arange(0, cur_num_samples), cur_roi_gt_labels] # 計算分類網絡和回歸網絡的損失值 cur_roi_cls_loss, cur_roi_reg_loss, _ = multitask_loss(out_offsets=cur_roi_offsets, out_classifier=cur_roi_out_classifier, gt_offsets=cur_roi_gt_offsets, gt_labels=cur_roi_gt_labels) roi_cls_loss += cur_roi_cls_loss roi_reg_loss += cur_roi_reg_loss # 計算整體loss, 反向傳播 batch_loss = (rpn_cls_loss + rpn_reg_loss + roi_cls_loss + roi_reg_loss) / num optimizer.zero_grad() batch_loss.backward() optimizer.step() # 記錄每輪訓練數據量和總loss train_loss += batch_loss.item() * num num_train_loss += num else: # 設置為驗證模式 network.eval() with torch.no_grad(): # 獲取輸入圖像全圖特征, 維度為[num, 512, im_height/16, im_width/16] im_features = network.forward(x=ims, mode="extractor") # 利用rpn網絡獲取回歸器輸出/分類器輸出/batch數據對應建議框/建議框對應數據索引/全圖先驗框 rpn_offsets, rpn_classifier, rois, rois_idx, anchors = network.forward(x=[im_features, (im_height, im_width)], mode="rpn") np_anchors = anchors.cpu().numpy() # 記錄rpn區域網絡的分類/回歸損失, 以及roi區域分類和回歸損失 rpn_cls_loss, rpn_reg_loss, roi_cls_loss, roi_reg_loss = 0, 0, 0, 0 samples_rois, samples_indexes, samples_offsets, samples_labels = [], [], [], [] # 遍歷每一個數據的真值框/數據標簽/rpn網絡輸出 for i in range(num): # 獲取每張圖像的真值框/標簽/rpn生成的偏移值/rpn生成的分類輸出 cur_gt_boxes = gt_boxes[i] cur_labels = labels[i] cur_rpn_offsets = rpn_offsets[i] cur_rpn_classifier = rpn_classifier[i] cur_rois = rois[i] np_cur_gt_boxes = cur_gt_boxes.clone().detach().cpu().numpy() np_cur_rois = cur_rois.clone().detach().cpu().numpy() np_cur_labels = cur_labels.clone().detach().cpu().numpy() # 根據當前圖像真值框和先驗框, 獲取每個先驗框標簽以及圖像經過rpn網絡后應產生的偏移值作為真值 cur_gt_rpn_labels, cur_gt_rpn_offsets = anchor_creator(im_width=im_width, im_height=im_height, anchors=np_anchors, gt_boxes=np_cur_gt_boxes) # 轉換為張量后計算rpn網絡的回歸損失和分類損失 cur_gt_rpn_offsets = torch.tensor(cur_gt_rpn_offsets).type_as(rpn_offsets) cur_gt_rpn_labels = torch.tensor(cur_gt_rpn_labels).long().to(rpn_offsets.device) cur_rpn_cls_loss, cur_rpn_reg_loss, _ = multitask_loss(out_offsets=cur_rpn_offsets, out_classifier=cur_rpn_classifier, gt_offsets=cur_gt_rpn_offsets, gt_labels=cur_gt_rpn_labels) rpn_cls_loss += cur_rpn_cls_loss rpn_reg_loss += cur_rpn_reg_loss # 在當前圖像生成的roi建議框中中, 抽取一定數量的正負樣本, 并計算出相應位置偏移, 用于后續回歸和區域分類 sample_rois, sample_offsets, sample_labels = proposal_creator(rois=np_cur_rois, gt_boxes=np_cur_gt_boxes, labels=np_cur_labels) # 將每個圖像生成的樣本信息存儲起來用于后續回歸和分類 # 抽取當前數據生成的推薦區域樣本放入list中, list長度為num, 每個元素維度為[num_samples, 4] samples_rois.append(torch.tensor(sample_rois).type_as(rpn_offsets)) # 將抽取的樣本索引放入list中, list長度為num, 每個元素維度為[num_samples] samples_indexes.append(torch.ones(len(sample_rois), device=rpn_offsets.device) * rois_idx[i][0]) # 將抽取的樣本偏移值放入list中, list長度為num, 每個元素維度為[num_samples, 4] samples_offsets.append(torch.tensor(sample_offsets).type_as(rpn_offsets)) # 將抽取的樣本分類標簽放入list中, list長度為num, 每個元素維度為[num_samples] samples_labels.append(torch.tensor(sample_labels, device=rpn_offsets.device).long()) # 整合當前batch數據抽取的推薦區域信息 samples_rois = torch.stack(samples_rois, dim=0) samples_indexes = torch.stack(samples_indexes, dim=0) # 將圖像特征和推薦區域送入模型, 進行roi池化并獲得分類模型/回歸模型輸出 roi_out_offsets, roi_out_classifier = network.forward(x=[im_features, samples_rois, samples_indexes, (im_height, im_width)], mode="head") for i in range(num): cur_num_samples = roi_out_offsets.size(1) cur_roi_out_offsets = roi_out_offsets[i] cur_roi_out_classifier = roi_out_classifier[i] cur_roi_gt_offsets = samples_offsets[i] cur_roi_gt_labels = samples_labels[i] cur_roi_out_offsets = cur_roi_out_offsets.view(cur_num_samples, -1, 4) # 根據roi對應的樣本標簽, 選擇與其類別對應真實框的offsets cur_roi_offsets = cur_roi_out_offsets[torch.arange(0, cur_num_samples), cur_roi_gt_labels] # 計算分類網絡和回歸網絡的損失值 cur_roi_cls_loss, cur_roi_reg_loss, _ = multitask_loss(out_offsets=cur_roi_offsets, out_classifier=cur_roi_out_classifier, gt_offsets=cur_roi_gt_offsets, gt_labels=cur_roi_gt_labels) roi_cls_loss += cur_roi_cls_loss roi_reg_loss += cur_roi_reg_loss # 計算整體loss, 反向傳播 batch_loss = (rpn_cls_loss + rpn_reg_loss + roi_cls_loss + roi_reg_loss) / num # 記錄每輪訓練數據量和總loss val_loss += batch_loss.item() * num num_val_loss += num scheduler.step() # 記錄loss和acc變化曲線 train_loss_all.append(train_loss / num_train_loss) val_loss_all.append(val_loss / num_val_loss) print("Epoch:[{:0>3}|{}] train_loss:{:.3f} val_loss:{:.3f}".format(epoch + 1, num_epochs, train_loss_all[-1], val_loss_all[-1])) # 保存模型 if val_loss_all[-1] < best_loss: best_loss = val_loss_all[-1] save_path = os.path.join("./model", "model_" + str(epoch + 1) + ".pth") torch.save(network, save_path) # 繪制訓練曲線 fig_path = os.path.join("./model/", "train_curve.png") plt.plot(range(num_epochs), train_loss_all, "r-", label="train") plt.plot(range(num_epochs), val_loss_all, "b-", label="val") plt.title("Loss") plt.legend() plt.tight_layout() plt.savefig(fig_path) plt.close() return None if __name__ == "__main__": device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = FasterRCNN(num_classes=CLASSES, train_flag=True, feature_stride=FEATURE_STRIDE, anchor_spatial_scales=ANCHOR_SPATIAL_SCALES, wh_ratios=ANCHOR_WH_RATIOS, pretrained=False) optimizer = torch.optim.Adam(params=model.parameters(), lr=LR) scheduler = StepLR(optimizer, step_size=STEP, gamma=GAMMA) model_root = "./model" os.makedirs(model_root, exist_ok=True) # 在生成的ss數據上進行預訓練 train_root = "./data/source" train_set = GenDataSet(root=train_root, im_width=IM_SIZE, im_height=IM_SIZE, batch_size=BATCH_SIZE, shuffle=True) train(data_set=train_set, network=model, num_epochs=EPOCHS, optimizer=optimizer, scheduler=scheduler, device=device, train_rate=0.8)
五、模型預測
1. 預測代碼
預測流程比訓練簡單的多，其具體過程如下：
- 將圖片輸入模型獲得Classifier網絡的分類輸出rois_out_classifier、偏移值輸出rois_out_regressor以及篩選后的推薦區域rois；
- 根據偏移值rois_out_regressor和rois計算預測的目標邊界框target_boxes；
- 利用rois_out_classifier分類結果對預測邊界框進行限制；
- 利用非極大值抑制nms對最終bbox進行篩選。
預測過程Predict.py代碼如下：
import os import torch import numpy as np from torch import Tensor import torch.nn.functional as F from torchvision.ops import nms from Model import FasterRCNN, ProposalCreator import matplotlib.pyplot as plt from matplotlib import patches from Utils import GenDataSet from skimage import io from Config import * def draw_box(img: np.ndarray, boxes: np.ndarray = None, save_name: str = None): """ 在圖像上繪制邊界框 :param img: 輸入圖像 :param boxes: bbox坐標, 列分別為[x, y, w, h, score, label] :param save_name: 保存bbox圖像名稱, None-不保存 :return: None """ plt.imshow(img) axis = plt.gca() if boxes is not None: for box in boxes: x, y, w, h = box[:4].astype("int") score = box[4] rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none') axis.add_patch(rect) axis.text(x, y - 10, "Score: {:.2f}".format(score), fontsize=12, color='blue') if save_name is not None: os.makedirs("./predict", exist_ok=True) plt.savefig("./predict/" + save_name + ".jpg") plt.show() return None def predict(network: FasterRCNN, im: np.ndarray, device: torch.device, im_width: int, im_height: int, num_classes: int, offsets_norm_params: Tensor, nms_thresh: float = 0.3, confidence_thresh: float = 0.5, save_name: str = None): """ 模型預測 :param network: Faster R-CNN模型結構 :param im: 原始輸入圖像矩陣 :param device: CPU/GPU :param im_width: 輸入模型的圖像寬度 :param im_height: 輸入模型的圖像高度 :param num_classes: 目標類別數 :param offsets_norm_params: 偏移值歸一化參數 :param nms_thresh: 非極大值抑制閾值 :param confidence_thresh: 目標置信度閾值 :param save_name: 保存文件名 :return: None """ # 測試模式 network.eval() src_height, src_width = im.shape[:2] # 數據歸一化和縮放 im_norm = (im / 255.0).astype("float32") im_rsz = GenDataSet.resize(im=im_norm, im_width=im_width, im_height=im_height, gt_boxes=None) # 將矩陣轉換為張量 im_tensor = torch.tensor(np.transpose(im_rsz, (2, 0, 1))).unsqueeze(0).to(device) with torch.no_grad(): # 獲取Faster R-CNN網絡的輸出, 包括回歸器輸出/分類器輸出/推薦區域 # 維度分別為[num_ims, num_rois, num_classes * 4]/[num_ims, num_rois, num_classes]/[num_ims, num_rois, 4] rois_out_regressor, rois_out_classifier, rois, _ = network.forward(x=im_tensor, mode="forward") # 獲取當前圖像數量/推薦區域數量 num_ims, num_rois, _ = rois_out_regressor.size() # 記錄預測的邊界框信息 out_bboxes = [] # 遍歷處理每張圖片, 此處實際只有一張圖 cur_rois_offsets = rois_out_regressor[0] cur_rois = rois[0] cur_rois_classifier = rois_out_classifier[0] # 將偏移值進行維度變換[num_rois, num_classes * 4] -> [num_rois, num_classes, 4] cur_rois_offsets = cur_rois_offsets.view(-1, num_classes, 4) # 對roi區域進行維度變換[num_rois, 4] -> [num_rois, 1, 4] -> [num_rois, num_classes, 4] cur_rois = cur_rois.view(-1, 1, 4).expand_as(cur_rois_offsets) # 將偏移值和roi區域展開成相同大小的二維張量, 方便對roi進行位置矯正 # 將偏移值展開, 維度[num_rois, num_classes, 4] -> [num_rois * num_classes, 4] cur_rois_offsets = cur_rois_offsets.view(-1, 4) # 將和roi區域展開, 維度[num_rois, num_classes, 4] -> [num_rois * num_classes, 4] cur_rois = cur_rois.contiguous().view(-1, 4) # 對回歸結果進行修正 # 注意Faster R-CNN網絡輸出的樣本偏移值是經過均值方差修正的, 此處需要將其還原 offsets_norm_params = offsets_norm_params.type_as(cur_rois_offsets) # ProposalTargetCreator中計算方式: sample_offsets = (sample_offsets - self.offsets_normalize_params[0]) / self.offsets_normalize_params[1] cur_rois_offsets = cur_rois_offsets * offsets_norm_params[1] + offsets_norm_params[0] # 利用偏移值對推薦區域位置進行矯正, 獲得預測框位置, 維度[num_rois * num_classes, 4] cur_target_boxes = ProposalCreator.calc_bboxes_from_offsets(offsets=cur_rois_offsets, anchors=cur_rois) # 展開成[num_rois, num_classes, 4]方便與類別一一對應 cur_target_boxes = cur_target_boxes.view(-1, num_classes, 4) # 獲取分類得分 cur_roi_scores = F.softmax(cur_rois_classifier, dim=-1) # 根據目標得分, 獲取最可能的分類結果 max_prob_labels = torch.argmax(input=cur_roi_scores[:, 1:], dim=-1) + 1 max_prob_scores = cur_roi_scores[torch.arange(0, cur_roi_scores.size(0)), max_prob_labels] max_prob_boxes = cur_target_boxes[torch.arange(0, cur_target_boxes.size(0)), max_prob_labels] # 選取得分大于閾值的 is_valid_scores = max_prob_scores > confidence_thresh if sum(is_valid_scores) > 0: valid_boxes = max_prob_boxes[is_valid_scores] valid_scores = max_prob_scores[is_valid_scores] keep = nms(boxes=valid_boxes, scores=valid_scores, iou_threshold=nms_thresh) # 獲取保留的目標框, 維度為[num_keep, 4] keep_boxes = valid_boxes[keep] # 獲取保留目標框的得分, 并將維度擴展為[num_keep, 1] keep_scores = valid_scores[keep][:, None] # 將預測框/標簽/得分堆疊在一起 cls_predict = torch.cat([keep_boxes, keep_scores], dim=1).cpu().numpy() # 預測結果添加進cur_out_bboxes里 out_bboxes.extend(cls_predict) if len(out_bboxes) > 0: out_bboxes = np.array(out_bboxes) # 計算原始輸入圖像和模型輸入圖像之間的空間縮放比例 map_scale = np.array([src_width, src_height, src_width, src_height]) / np.array([im_width, im_height, im_width, im_height]) # 將預測框從模型輸入圖像映射到原始輸入圖像 out_bboxes[:, :4] = out_bboxes[:, :4] * map_scale # 對預測框坐標進行限制 out_bboxes[:, [0, 2]] = np.clip(a=out_bboxes[:, [0, 2]], a_min=0, a_max=src_width - 1) out_bboxes[:, [1, 3]] = np.clip(a=out_bboxes[:, [1, 3]], a_min=0, a_max=src_height - 1) # 將預測框[x1, y1, x2, y2, score, label]轉換為[x1, y1, w, h, score, label] out_bboxes[:, [2, 3]] = out_bboxes[:, [2, 3]] - out_bboxes[:, [0, 1]] + 1 if len(out_bboxes) == 0: out_bboxes = None draw_box(img=im, boxes=out_bboxes, save_name=save_name) if __name__ == "__main__": device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model_path = "./model/model_180.pth" model = torch.load(model_path, map_location=device) # 偏移值歸一化參數, 需保證和訓練階段一致 offsets_normalize_params = torch.Tensor([[0, 0, 0, 0], [0.1, 0.1, 0.2, 0.2]]) test_root = "./data/source/17flowers" for roots, dirs, files in os.walk(test_root): for file in files: if not file.endswith(".jpg"): continue im_name = file.split(".")[0] im_path = os.path.join(roots, file) im = io.imread(im_path) predict(network=model, im=im, device=device, im_width=IM_SIZE, im_height=IM_SIZE, num_classes=CLASSES, offsets_norm_params=offsets_normalize_params, nms_thresh=NMS_THRESH, confidence_thresh=CONFIDENCE_THRESH, save_name=im_name)
2. 預測結果
如下為Faster RCNN在花朵數據集上預測結果展示，左圖中當同一圖中存在多個相同類別目標且距離較近時，邊界框并沒有很好地檢測出每個個體。
六、算法缺點
Faster R-CNN雖然在目標檢測領域取得了顯著成果，但仍存在一些缺點：
- 卷積提取網絡的問題：Faster R-CNN在特征提取階段，無論使用VGGNet還是ResNet，其特征圖都是單層的，且分辨率通常較小。這可能導致對于多尺度、小目標的檢測效果不佳。為了優化這一點，研究者們提出了使用多層融合的特征圖或增大特征圖的分辨率等方法；
- NMS（非極大值抑制）的問題：NMS在RPN產生Proposal時用于避免重疊的框，并以分類得分為篩選標準。但NMS對于遮擋物體的處理并不友好，有可能將本屬于兩個物體的Proposal過濾為一個，導致漏檢。因此，對NMS的改進可以進一步提升檢測性能；
- RoI Pooling的問題：Faster R-CNN的原始RoI Pooling在兩次取整過程中可能會帶來精度的損失，雖然后續的Mask R-CNN針對此問題進行了改進，但原始的Faster R-CNN在這一方面仍有待優化。
七、數據和代碼
本文中數據和詳細工程代碼實現請移步：https://github.com/jchsun1/Faster-RCNN
本文至此結束，如有疑惑歡迎留言交流。
