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くらい出るようになり、
待ち時間を大幅に減らすことができました。

長くなったので今回はここまで。
次は学習処理について書こうと思います。