Google Colaboratory + pytorch で SIGNATEコンペ参加(開発編3)
引き続きGoogle Colaboratory + pytorch で SIGNATEの
AIエッジコンテスト(オブジェクト検出)に参加するお話です。
前回までで、モデル、Loss関数など学習に必要な部分を揃えました。
後は実際にオブジェクト検出を行い、正解と比較して評価を行う部分を作れば完成です。
オブジェクト検出処理
今回オブジェクト検出処理用にpredictという関数を用意しました。
この関数ではネットワークにより計算されたlocとclassを矩形に変換し、
リストにして出力します。
ただし、batchsizeは1を前提としています。
def predict(image, score_th=0.7, nms_th=0.3, class_agnostic=False): image = image.to(device) pred_loc,pred_cls,anchor = net(image) anchor = anchor.to(device) loc_normalize_std = torch.tensor([0.1,0.1,0.2,0.2]).to(device) pred_cls = torch.nn.functional.softmax(pred_cls,2) pred_loc_batch = pred_loc[0] detect_box_list = [] detect_score_list = [] detect_label_list = [] num_class = pred_cls.size(2) - 1 _classes = ('Car','Truck','Pedestrian', 'Bicycle','Signal','Signs') for class_idx in range(num_class): if(class_agnostic): pred_cls_loc = pred_loc_batch.detach() else: pred_cls_loc = pred_loc_batch[:,class_idx*4:(class_idx+1)*4].detach() pred_cls_loc = pred_cls_loc * loc_normalize_std if _classes[class_idx] in size_expansion: sw = 1./size_expansion[_classes[class_idx]][0] sh = 1./size_expansion[_classes[class_idx]][1] pred_cls_loc[:,2:] = pred_cls_loc[:,2:] * torch.tensor([sw, sh]).to(device) pred_box = loc2box(pred_cls_loc.detach(),anchor.detach()) detect_idx = (pred_cls[0][:,class_idx+1] >= score_th) detect_box = pred_box[detect_idx] detect_score = pred_cls[0][detect_idx,class_idx+1] if(detect_box.size(0) > 0): detect_box[:,0::2].clamp_(min=0,max=image.size(3)-1) detect_box[:,1::2].clamp_(min=0,max=image.size(2)-1) keep = box_nms(detect_box,detect_score,threshold=nms_th) detect_box_list.append(detect_box[keep]) detect_score_list.append(detect_score[keep]) detect_label_list.append(torch.Tensor(keep.size(0)).fill_(class_idx).to(device)) if(len(detect_box_list)): detect_box_list = torch.cat(detect_box_list,0) detect_score_list = torch.cat(detect_score_list,0) detect_label_list = torch.cat(detect_label_list,0) return detect_box_list, detect_score_list, detect_label_list
そしてこの関数で得られた結果を、今回のコンペで評価できる形に変換します。
result_data = OrderedDict() _classes = ('Car','Truck','Pedestrian', 'Bicycle','Signal','Signs') _, order = torch.sort(detect_score_list, dim=0, descending=True) order_detect_box_list = ((detect_box_list * scale + 0.5).int())[order] order_detect_label_list = detect_label_list[order] for i in range(order.size(0)): idx = order[i] label_name = _classes[int(detect_label_list[idx])] if(not (label_name in result_data)): result_data[label_name] = [] if(len(result_data[label_name]) < 100): x1 = out_detect_box_list[idx,0] y1 = out_detect_box_list[idx,1] x2 = out_detect_box_list[idx,2] y2 = out_detect_box_list[idx,3] result_data[label_name].append([x1,y1,x2,y2])
mAPの評価関数等はSIGNATEであらかじめ用意されていたものをそのまま使いました。
以上で学習されたモデルを使って、評価を行うことができました。
実際に出来上がったソースを使って一旦投稿してみたところ、
当然ですが、最下位に近い成績でした。。。
mAP 0.07位だったと思います。
ただ、この時点でコンペに参加するという目標は達成できました。
オブジェクト検出処理の改善
さて、上にあげた検出処理ですが2つの点で問題があり、検出処理に5秒/1フレームほどかかってしまい、
学習の大半がTESTの時間になってしまうという事態になりました。
そこで、少しだけ処理速度の改善を行ったので紹介いたします。
処理結果の変換処理
上で上げたソースでは、for文を使ってデータ形式を変換していますが、
pythonのfor文はとても遅く、ボトルネックとなっていました。
そこで、以下の様に変更し、pytorchにできるだけ計算を頑張ってもらうようにし
処理速度を上げました。
result_data = OrderedDict() _classes = ('Car','Truck','Pedestrian', 'Bicycle','Signal','Signs') _, order = torch.sort(detect_score_list, dim=0, descending=True) order_detect_box_list = ((detect_box_list * scale + 0.5).int())[order] order_detect_label_list = detect_label_list[order] for label_idx, label_name in enumerate(_classes): label_ids = (order_detect_label_list == label_idx) label_box_list = order_detect_box_list[label_ids][:100].cpu().numpy().tolist() if len(label_box_list) > 0: result_data[label_name] = label_box_list
Non Maximum Suppression
上のソースではさらっと無視していましたが、
重なって検出された矩形たちをまとめる処理としてNon Maximum Suppressionを使っています。
ここでは、以下の様なCPUベースの処理を使っています。
この実装はpytorchのSSDやRetinanetの実装などにあったものをそのままお借りしています。
def box_nms(bboxes, scores, threshold=0.3, mode='union'): '''Non maximum suppression. Args: bboxes: (tensor) bounding boxes, sized [N,4]. scores: (tensor) bbox scores, sized [N,]. threshold: (float) overlap threshold. mode: (str) 'union' or 'min'. Returns: keep: (tensor) selected indices. Reference: https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/nms/py_cpu_nms.py ''' x1 = bboxes[:,0] y1 = bboxes[:,1] x2 = bboxes[:,2] y2 = bboxes[:,3] areas = (x2-x1+1) * (y2-y1+1) _, order = scores.sort(0, descending=True) keep = [] while order.numel() > 0: i = order[0] keep.append(i) if order.numel() == 1: break xx1 = x1[order[1:]].clamp(min=x1[i]) yy1 = y1[order[1:]].clamp(min=y1[i]) xx2 = x2[order[1:]].clamp(max=x2[i]) yy2 = y2[order[1:]].clamp(max=y2[i]) w = (xx2-xx1+1).clamp(min=0) h = (yy2-yy1+1).clamp(min=0) inter = w*h if mode == 'union': ovr = inter / (areas[i] + areas[order[1:]] - inter) elif mode == 'min': ovr = inter / areas[order[1:]].clamp(max=areas[i]) else: raise TypeError('Unknown nms mode: %s.' % mode) ids = (ovr<=threshold).nonzero().view(-1) if ids.numel() == 0: break order = order[ids+1] return torch.LongTensor(keep)
今回はGoogle Colaboratoryで処理を行うので、CUDA周りの設定はいろいろ面倒と考え、
これでよいだろうと思っていたのですが、やはり無視できない遅さでしたので、
cupyベースのNMSを使うことにしました。
ソースは以下からお借りしました。
github.com
nmsの部分をあらかじめGoogle Driveの中に仕込んで置き、
以下を実行することで、ビルドを行います。
#setup nms_cupy !mkdir -p model/utils/ !cp -r '/content/gdrive/My Drive/colab_pytorch_detection/nms' /content/model/utils/ %cd /content/model/utils/nms/ !python build.py build_ext --inplace %cd /content
build.pyは元のままだとビルドできないので、
numpy周りの参照関係を修正しておきます。
そして、cupy用のbox_nms関数を用意します。
import cupy as cp from model.utils.nms import non_maximum_suppression def box_nms_cupy(bboxes, scores, threshold=0.3): boxes_np = bboxes.detach().cpu().numpy() prob_np = scores.detach().cpu().numpy() keep = non_maximum_suppression( cp.array(boxes_np), threshold, prob_np) keep = torch.tensor(cp.asnumpy(keep)).long() return keep
これを最初のソースのbox_nmsの代わりに使えばOKです。
以上2点を直すことで、処理速度が5fpsくらい出るようになり、
待ち時間を大幅に減らすことができました。
長くなったので今回はここまで。
次は学習処理について書こうと思います。