SSDのハードネガティブマイニングの結果を可視化してみる

ハードネガティブマイニングでネガティブデフォルトボックスから選ばれた、ラベル付けの損失に寄与するネガティブデフォルトボックスを可視化させてみたいと思います。実装はSSDのPython実装amdegroot/ssd.pytorchを使用します。

『SSDのポジティブDBoxを表示してみる』のコードの再掲になりますが、以下のようにしてjaccard係数が最大になるバウンディングボックスとデフォルトボックスとの位置情報をloc_tに、そのバウンディングボックスに対応するラベル情報をconf_tにデフォルトボックスのインデックスと合うように入れます。loc_tをmatch関数の定義の関係上で定義していますが、この変数は今回使用されることはありません。

image_size = 300
batch_size = 1
dataset = VOCDetection(root = '/path/to/root', 
                          transform=BaseTransform(image_size, MEANS))
data_loader = data.DataLoader(dataset, batch_size = batch_size, num_workers=0, shuffle=False, collate_fn = detection_collate)
torch.set_default_tensor_type('torch.cuda.FloatTensor')
images, targets = next(iter(data_loader))
images = images.cuda()
targets = [ann.cuda() for ann in targets]
priors = PriorBox(voc).forward()
num_priors = priors.size(0)
loc_t = torch.Tensor(batch_size, num_priors,  priors.size(1))
conf_t = torch.LongTensor(batch_size, num_priors)
match(threshold = 0.5, 
      truths = targets[0][:, :-1].data,  
      priors = priors.data, 
      variances = [0.1, 0.2], 
      labels = targets[0][:, -1].data, 
      loc_t = loc_t, 
      conf_t = conf_t, 
      idx = 0)

ラベル付けの学習では、直感的に背景を背景として判定できるようにネガティブデフォルトボックスの損失情報も取り入れる必要があることは分かると思いますが、SSDではネガティブデフォルトボックスのうち、うまくラベル付けできていないものだけを選ぶことをします。この操作をハードネガティブマイングと呼びます。

ハードネガティブマイニングをするには、モデルにラベル付けする必要があるので、SSDをトレーニングモードで生成して画像のラベル付け(conf_data)を計算させます。

net = build_ssd('train', 300, len(VOC_CLASSES) + 1)
# init.pthはtrain.pyの学習開始時の重み
net.load_state_dict(torch.load("/path/to/init.pth"))
net.vgg.load_state_dict(torch.load(f"{os.getcwd()}/weights/vgg16_reducedfc.pth"))
net = net.cuda()
net.train()
_, conf_data, _ = net(images)

conf_dataの計算で使用するSSDの重みは、原理的にはどのような重みでも良いのですが、ここでは学習開始時の重みを使用することにします。

以下のコードが所謂ハードネガティブマイニングにあたるものですが、ラベル付けの教師データconf_tに対してモデルの予測conf_dataがどれだけズレているかをクロスエントロピーで評価し、ソートして損失の大きいもの上位をとっていくことをしています。

#ポジティブデフォルトボックスに対して選ぶネガティブデフォルトボックスの個数が何倍あるか
#学習処理では3に設定されていましたが、最後画像として描画する都合のために2にしています
dbox_ratio = 2
conf_t = Variable(conf_t.cuda())
pos = conf_t > 1
batch_conf = conf_data.view(-1, 21)
loss_c = F.cross_entropy(batch_conf, conf_t.view(-1), reduction='none')
num_pos = pos.long().sum(1, keepdim = True)
loss_c = loss_c.view(1, -1)
loss_c[pos] = 0
_, loss_idx = loss_c.sort(1, descending = True)
_, idx_rank = loss_idx.sort(1) 
num_neg = torch.clamp(dbox_ratio * num_pos, max = pos.size(1) - 1)
neg = idx_rank < num_neg.expand_as(idx_rank)

デフォルトボックスの内、negがTrueになっているインデックスの要素が選ばれたネガティブデフォルトボックスになっています。これらを以下のコードで画像として表示してみます。

image = (images[0].to('cpu').detach().numpy().transpose(1, 2, 0) + (MEANS[2], MEANS[1], MEANS[0])).astype(np.uint8).copy()   
image = cv2.resize(image, (image_size, image_size))
plt.figure(figsize=(80, 180))
indices = [i for i, v in enumerate(list(neg.to('cpu').detach().numpy().copy()[0])) if v]
for i, idx in enumerate(indices):
    img = image.copy()
    cx_d, cy_d, w_d, h_d = priors[idx].to('cpu').detach().numpy().copy()
    xmin_d = int((cx_d - w_d / 2) * image_size)
    ymin_d = int((cy_d - h_d / 2) * image_size)
    xmax_d = int((cx_d + w_d / 2) * image_size)
    ymax_d = int((cy_d + h_d / 2) * image_size)
    pt1_d = (xmin_d, ymin_d)
    pt2_d = (xmax_d, ymax_d)    
    cv2.rectangle(img, pt1=pt1_d, pt2=pt2_d, color=(0, 255, 0), thickness=2)
    for box in targets[0]:
        label = int(box[4])
        pt1 =  (int(box[0] * image_size), int(box[1] * image_size) )
        pt2 = (int(box[2] * image_size), int(box[3] * image_size))
        jaccard = calc_jaccard((pt1_d[0], pt1_d[1], pt2_d[0], pt2_d[1]), (pt1[0], pt1[1], pt2[0], pt2[1]))
        cv2.putText(img, f"{jaccard: .2f}", (pt1[0] + 5 , pt1[1] - 5), cv2.FONT_HERSHEY_PLAIN, 1.5, (0, 255, 0), 2)
        cv2.rectangle(img, pt1=pt1, pt2=pt2, color=(255, 0, 0), thickness=2)
    plt.subplot(9 , 4, i + 1)
    plt.axis('off')
    plt.imshow(img)
plt.show()

小さなデフォルトボックスが選ばれる傾向があるような気がしますが、なぜなのかは分かっていません。選ばれているデフォルトボックスが、バウンディングボックスに含まれるのは基底モデルとしてvgg16の学習済みモデルを使用しているため、dogなどの動物に重みが高めになるためクロスエントロピーが大きくなる傾向があるように思います。

最後に学習済みモデルで傾向が変わるか確認してみます。今回、学習済みモデルとして使用した重みは、次のURLからダウンロードできます。
https://s3.amazonaws.com/amdegroot-models/ssd300_mAP_77.43_v2.pth

結果は学習前よりも大きなデフォルトボックスも選ばれるようになり、dogと分類するのが自然と思えるものも含まれてるようになります。こうして見ると、SSDの定めているルールと画像分類vgg16との判定は、ネガティブデフォルトボックスにおいては常に対立しているように見えますね。

ここで使用している元画像はここ (Pixabay License: 商用利用無料 帰属表示は必要ありません)から640x426のサイズでダウンロードしたものになります。