z [코드 개인 공부] : Single Shot Multibox Detector (SSD)
본문 바로가기

Computer Vision

[코드 개인 공부] : Single Shot Multibox Detector (SSD)

728x90

개요

여러 블로그에서 SSD에 대한 설명은 충분히 나와있습니다만, 막상 코드에서 어떻게 돌아가는지를 파악할 때 어려움이 굉장히 많았어서 SSD가 실행 될 때 중요 부분 코드를 뜯어보고자 작성했습니다. (개인 공부 목적, Image Tracking 목적)

SSD에서 여러번 prediction하는 layer를 포함하는 부분 등 깃허브의 코드를 봤을 때 자명하게 알 수 있는 부분들은 제외했습니다.

github :  https://github.com/amdegroot/ssd.pytorch

 

 

SSD 과정 diagram

 

SSD Module

class SSD(nn.Module):

    def __init__(self, phase, size, base, extras, head, num_classes):
        super(SSD, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.cfg = (coco, voc)[num_classes == 21]
        self.priorbox = PriorBox(self.cfg) # Generate Anchor Box Class
        self.priors = Variable(self.priorbox.forward(), volatile=True) # Generate Anchor Box
        self.size = size

        # SSD network
        self.vgg = nn.ModuleList(base)
        # Layer learns to scale the l2 normalized features from conv4_3
        self.L2Norm = L2Norm(512, 20)
        self.extras = nn.ModuleList(extras)

        self.loc = nn.ModuleList(head[0])
        self.conf = nn.ModuleList(head[1])

        if phase == 'test':
            self.softmax = nn.Softmax(dim=-1)
            self.detect = Detect(num_classes, 0, 200, 0.01, 0.45)

    def forward(self, x):

        sources = list()
        loc = list()
        conf = list()

        # apply vgg up to conv4_3 relu
        for k in range(23):
            x = self.vgg[k](x)

        s = self.L2Norm(x)
        sources.append(s)

        # apply vgg up to fc7
        for k in range(23, len(self.vgg)):
            x = self.vgg[k](x)
        sources.append(x) # source에 예측을 진행할 실질적인 layers를 넣는다.

        # apply extra layers and cache source layer outputs
        for k, v in enumerate(self.extras):
            x = F.relu(v(x), inplace=True)
            if k % 2 == 1:
                sources.append(x)

        # apply multibox head to source layers
        for (x, l, c) in zip(sources, self.loc, self.conf): # self.loc : 각 feature maps의 위치에서 predict를 실행한다. ratio_constant * 4
            loc.append(l(x).permute(0, 2, 3, 1).contiguous())
            conf.append(c(x).permute(0, 2, 3, 1).contiguous())

        """
            img_size = 300, bs = 32
            torch.Size([32, 38, 38, 16]) -> torch.Size([32, 23104])
            torch.Size([32, 19, 19, 24]) -> torch.Size([32, 8664])
            torch.Size([32, 10, 10, 24]) -> torch.Size([32, 2400])
            torch.Size([32, 5, 5, 24]) -> torch.Size([32, 600])
            torch.Size([32, 3, 3, 16]) -> torch.Size([32, 144])
            torch.Size([32, 1, 1, 16]) -> torch.Size([32, 16])
            32, -1에서, -1끼리 concat 시킨다. concat(axis = 1)
            
            이 후 reshape (bs, -1, 4) -1 부분은 각 feature maps, each ratio에 대한 모든 prediction
            
        """

        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1) # bs, 나머지 concat
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
        if self.phase == "test":
            output = self.detect(
                # shape -1 에서 각 feature maps, each ratios 에 대한 결과가 모두 들어가있다.
                loc.view(loc.size(0), -1, 4),                   # loc preds
                self.softmax(conf.view(conf.size(0), -1,
                             self.num_classes)),                # conf preds
                self.priors.type(type(x.data))                  # default boxes
            )
        else:
            output = (
                loc.view(loc.size(0), -1, 4),
                conf.view(conf.size(0), -1, self.num_classes),
                self.priors # 각 loc, conf 결과에 대한 anchors
            )
        return output

이는 pytorch SSD Official 코드 입니다. https://github.com/amdegroot/ssd.pytorch

 

 

        for (x, l, c) in zip(sources, self.loc, self.conf): # self.loc : 각 feature maps의 위치에서 predict를 실행한다. ratio_constant * 4
            loc.append(l(x).permute(0, 2, 3, 1).contiguous())
            conf.append(c(x).permute(0, 2, 3, 1).contiguous())

SSD 모듈에서 순전파 부분의 한 부분입니다. 이는 각 feature maps (6개)의 각 위치에서와 각 ratio에서의 prediction 진행 부분입니다.

 

 

"""
            img_size = 300, bs = 32
            torch.Size([32, 38, 38, 16]) -> torch.Size([32, 23104])
            torch.Size([32, 19, 19, 24]) -> torch.Size([32, 8664])
            torch.Size([32, 10, 10, 24]) -> torch.Size([32, 2400])
            torch.Size([32, 5, 5, 24]) -> torch.Size([32, 600])
            torch.Size([32, 3, 3, 16]) -> torch.Size([32, 144])
            torch.Size([32, 1, 1, 16]) -> torch.Size([32, 16])
            32, -1에서, -1끼리 concat 시킨다. concat(axis = 1)
            
            이 후 reshape (bs, -1, 4) -1 부분은 각 feature maps, each ratio에 대한 모든 prediction
            
        """

        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1) # bs, 나머지 concat
        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
        if self.phase == "test":
            output = self.detect(
                # shape -1 에서 각 feature maps, each ratios 에 대한 결과가 모두 들어가있다.
                loc.view(loc.size(0), -1, 4),                   # loc preds
                self.softmax(conf.view(conf.size(0), -1,
                             self.num_classes)),                # conf preds
                self.priors.type(type(x.data))                  # default boxes
            )
        else:
            output = (
                loc.view(loc.size(0), -1, 4),
                conf.view(conf.size(0), -1, self.num_classes),
                self.priors # 각 loc, conf 결과에 대한 anchors
            )

ratios constants : default box의 사전 ratio 설정에 따라 예측의 수가 달라집니다. 논문 상에서는 2:1 이런식으로 ratio를 지정하면

(자기 자신, 1:1로 확대된 자기자신, 2:1 가로세로 각각) 이렇게 4가지가 나옵니다.

ratios constants가 6으로 나온 경우는 ratio의 종류가 2가지라는 의미입니다. (2:1, 3:1 이런식으로 ..)

 

이제 여러 feature maps에서 (h*w*ratio constants) 만큼 예측을 시도합니다. 이에 대해서 (Batch size, -1, 4 or num_classes)로 reshape 함으로써 처리를 편하게 합니다.

-1이 의미하는 바는 각 feature maps에서 각 위치와 각 ratios에 대한 prediction 수 입니다. (6 * h * w * ratios constants)

 

Output은 위의 reshape한 형태로 내보냅니다.

 

 

Multibox Loss (Train)

default box를 미리 지정해놓고 정답의 bounding box와 가장 비슷한 (iou가 가장 큰) default box를 매칭하여 train에 사용되는 정답을 매칭된 default box에 대해 정답의 bounding box를 encode 합니다. (논문에 하는 법 나와있음.)

 

하기 나름이긴 한데 어떤 코드에서는 사전에 default box와 ground truth를 encode 해놓기도 하고 어떤 코드(오피셜)에서는 loss에서 매칭을 진행합니다.

 

class MultiBoxLoss(nn.Module):

    def __init__(self, num_classes, overlap_thresh, prior_for_matching,
                 bkg_label, neg_mining, neg_pos, neg_overlap, encode_target,
                 use_gpu=True):
        super(MultiBoxLoss, self).__init__()
        self.use_gpu = use_gpu
        self.num_classes = num_classes
        self.threshold = overlap_thresh
        self.background_label = bkg_label
        self.encode_target = encode_target
        self.use_prior_for_matching = prior_for_matching
        self.do_neg_mining = neg_mining
        self.negpos_ratio = neg_pos
        self.neg_overlap = neg_overlap
        self.variance = cfg['variance']

    def forward(self, predictions, targets):
        """Multibox Loss
        Args:
            predictions (tuple): A tuple containing loc preds, conf preds,
            and prior boxes from SSD net.
                conf shape: torch.size(batch_size,num_priors,num_classes)
                loc shape: torch.size(batch_size,num_priors,4)
                priors shape: torch.size(num_priors,4)

            targets (tensor): Ground truth boxes and labels for a batch,
                shape: [batch_size,num_objs,5] (last idx is the label).
        """
        loc_data, conf_data, priors = predictions
        num = loc_data.size(0)
        priors = priors[:loc_data.size(1), :] # 혹시라도 설정이 달라서 anchor box 와 다른 경
        num_priors = (priors.size(0))
        num_classes = self.num_classes

        # match priors (default boxes) and ground truth boxes
        loc_t = torch.Tensor(num, num_priors, 4)
        conf_t = torch.LongTensor(num, num_priors)
        for idx in range(num):
            truths = targets[idx][:, :-1].data # gt에서의 bbox
            labels = targets[idx][:, -1].data # gt에서의 label
            defaults = priors.data # default box
            match(self.threshold, truths, defaults, self.variance, labels,
                  loc_t, conf_t, idx)
            # output이 loc_t, conf_t 생각하면 될듯
        if self.use_gpu:
            loc_t = loc_t.cuda()
            conf_t = conf_t.cuda()
        # wrap targets
        loc_t = Variable(loc_t, requires_grad=False)
        conf_t = Variable(conf_t, requires_grad=False)

        pos = conf_t > 0 # 배경이 존재한다 생각 될 때만. (여기에서 regress 기준이 정해진다.)
        num_pos = pos.sum(dim=1, keepdim=True)

        # Localization Loss (Smooth L1)
        # Shape: [batch,num_priors,4]
        pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
        loc_p = loc_data[pos_idx].view(-1, 4)
        loc_t = loc_t[pos_idx].view(-1, 4)
        loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)

        # Compute max conf across batch for hard negative mining
        batch_conf = conf_data.view(-1, self.num_classes)
        loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))

        # Hard Negative Mining
        loss_c[pos] = 0  # filter out pos boxes for now
        loss_c = loss_c.view(num, -1)
        _, loss_idx = loss_c.sort(1, descending=True)
        _, idx_rank = loss_idx.sort(1)
        num_pos = pos.long().sum(1, keepdim=True)
        num_neg = torch.clamp(self.negpos_ratio*num_pos, max=pos.size(1)-1)
        neg = idx_rank < num_neg.expand_as(idx_rank)

        # Confidence Loss Including Positive and Negative Examples
        pos_idx = pos.unsqueeze(2).expand_as(conf_data)
        neg_idx = neg.unsqueeze(2).expand_as(conf_data)
        conf_p = conf_data[(pos_idx+neg_idx).gt(0)].view(-1, self.num_classes)
        targets_weighted = conf_t[(pos+neg).gt(0)]
        loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)

        # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N

        N = num_pos.data.sum()
        loss_l /= N
        loss_c /= N
        return loss_l, loss_c
# match priors (default boxes) and ground truth boxes
        loc_t = torch.Tensor(num, num_priors, 4)
        conf_t = torch.LongTensor(num, num_priors)
        for idx in range(num):
            truths = targets[idx][:, :-1].data # gt에서의 bbox
            labels = targets[idx][:, -1].data # gt에서의 label
            defaults = priors.data # default box
            match(self.threshold, truths, defaults, self.variance, labels,
                  loc_t, conf_t, idx)
            # output이 loc_t, conf_t 생각하면 될듯

위 multibox loss의 코드에서 이 부분이 matching하는 부분입니다. truths와 default box를 매칭시켜 encode까지 한 후 내보냅니다.

official https://github.com/amdegroot/ssd.pytorch/blob/master/layers/box_utils.py 의 match 및 encode method를 살펴보면 이해가 갑니다.

 

위의 loss를 가지고 훈련이 진행됩니다.

 

Test, 이후 SMOT 생각.

훈련이 진행 될 때 각 default boxes에 대해 encode 된 data(정답)을 가지고 훈련이 되었습니다.
이런 형식 (각 default boxes의 regress 형식)으로 훈련이 되었기 때문에 evaluation 때도 image에 대해 해당 default box의 regression으로 진행됩니다. 이 prediction 값을 decode 하는 과정이 필요합니다. (여러 pred가 나오기 때문에 nms 까지)
각 feature maps의 each ratios, scales의 anchor boxes에 regress 진행됩니다.
SMOT(Single Shot Multi Object Tracking)에서는 anchor assignment시 output의 index 지정이 중요하지 않을까 싶습니다.

SSD 에서는 output(loc)이 (Batch Size, -1, 4) 형식이다. -1은 각 feature maps의 h*w*ratios를 모두 합친 것. (38, 19, 10, 5, 1)의 순서로 예측을 진행하고, prior (anchor box) 또한 이 순서와 맞습니다.
priors의 shape도 (Batch Size, -1, 4)형식. (SMOT 시 anchor box의 index를 그대로 지정해 줄 수 있음을 의미합니다.)

SMOT 때는 할당한 anchor indices (iou table에 따른 Hungarian Algorithm 등의 Match algorithm 사용)를 구합니다.
이 anchor indices 를 받고 output에서의 anchor box index로 아웃풋을 가져오면 될 것 같습니다. (여러 ratios가 존재할 거 같은데 detect(nms)로 해결하면 될 것 같다.)

 

 

728x90