[최적화] Structural Pruning via Latency-Saliency Knapsack 논문 설명 및 구현

들어가기 전에...

https://arxiv.org/pdf/2210.06659.pdf

https://github.com/NVlabs/HALP

2022년

원래는 모델 경량화의 초기 논문부터 쭉 보려고 했으나, 현재 연구 방향을 기준으로 관련 개념들을 공부하는 것이 더 유익할 듯하여, 2020년 이후 논문 중 github가 존재하는 논문을 보기로 했다. 단순히 논문의 text를 읽는 것이 중요한 것이 아니라, 코드도 같이 보면서 어떻게 구현하였는지 까지 보는 게 좋을 것 같다. 아마도 이해가 안되는 부분이 매우 많을 것으로 예상되나,, 열심히 찾아보면서 도전하는거지!

그리고, 논문의 포맷이나 순서대로 정리하기보다는 내가 공부한 순서대로 글을 적어보자

파랑색 글씨는 내가 보충한 내용 + 생각한 내용

빨강색 글씨는 아직 이해중인 내용

 

 

Paper - Abstract

 

구조적 가지치기는 랜덤하게 가지치기를 하는 비구조적 가지치기 방식과 달리 그룹(ex. channel, filter, layer, etc) 단위로 가지치기를 수행하는 것이다. 이는 구조 자체를 제거함으로써 네트워크 아키텍처를 단순화하여 inference 속도를 향상시킬 수 있다. 본 논문에서는 구조적 가지치기를 global resource allocation 최적화 문제로 공식화하는 HALP(Hardware-Aware Latency Pruning)를 제안한다. 이는 targeting device에 대해 미리 정의된 limit에 따라 latency를 제한하면서 정확도를 최대화하는 것을 목표로 한다. 

filter 중요도 순위를 위해, HALP는 latency lookup table을 활용하여 latency reduction potential과 global saliency score를 추적하여 정확도 하락을 측정한다. 두 metric(latency reduction potential과 global saliency score) 모두 가지치기 중에 효율적으로 평가할 수 있으므로 target constraint에서 reward maximization problem 하에서 global 구조적 가지치기를 재구성할 수 있다. 이를 통해 our augmented knapsack solver를 통해 (knapsack 알고리즘!?: 최대 용량이 존재하는 배낭에 가치와 무게가 다른 물건을 담아야할 때, 최대 용량을 넘지 않으며+가치를 최대로 하여 물건을 담는 방법을 찾는 알고리즘)HALP가 가지치기와 정확도의 효율성 trade-off에서 선행 작업을 능가할 수 있다. 다양한 네트워크 및 플랫폼에서 분류 및 탐지 작업에 대해 HALP를 검증하였다. 특히, ImageNet의 ResNet-50/-101 가지치기의 경우, HALP는 각각 +0.3%/-0.2% top-1 차이로 네트워크 처리량을 1.60배/1.90배 향상시킨다. VOC의 SSD 가지치기의 경우, HALP는 0.56mAP만의 감소로 처리량을 1.94배 향상시킨다. HALP는 선행 기술을 큰 차이를 보이며 능가한다.

 

project page : https://halp-neurips.github.io/

 

 

 

Github

 

github 접속시, 바로 네트워크 구조 그림을 확인할 수 있었다.

처음 이미지를 보고 해석한 바는 다음과 같다.

• 기본 network에서 filter 부분에 대해서 가지치기를 수행할 예정인듯하다
• 4개의 filter가 있을 때, 각각의 중요도를 추정하고, 높은 값을 가지는 2개의filter와, 낮은 값을 가지는 2개의 filter를 grouping한다. -> 이것이 accuracy guidance
• profiled lookup table을 기준으로? 앞서 grouping된 filter의 latency를 추정한다. -> 이것이 latency guidance
• 앞서 구한 accuracy, latency guidance를 통해 network global filter profile을 파악한다.
• 이를 활용하여 augmented knapsack을 수행한다.
• 이를 통해 나온 결과로 filter를 선택한다.

 

 

Project page

 

global resource allocation -> 3.1 설명

augmented knapsack 알고리즘 -> 3.2 설명

실험 결과 -> 4 설명

전체 컨셉은 knapsck 알고리즘을 쓴다는것을 기본으로 한다. filter의 정확도를 최대로 하며, cimpute cost를 활용하여 하드웨어 limit을 넘지 않도록 하는 filter를 선택하는 것이다.

가지치기에서 학습은 iterative때문에 존재하는 것인가?

 

 

Paper - Method

 

global resource allocation -> 3.1 설명 ( 가지치기 과정을 최적화 과정으로 공식화하여, auccuracy와 latency를 위한 중요도 추정을 설명)

augmented knapsack 알고리즘 -> 3.2 설명

실험 결과 -> 4 설명

전체 컨셉은 knapsck 알고리즘을 쓴다는것을 기본으로 한다. filter의 정확도를 최대로 하며, cimpute cost를 활용하여 하드웨어 limit을 넘지 않도록 하는 filter를 선택하는 것이다.

가지치기에서 학습은 iterative때문에 존재하는 것인가?

 

Objective Function

선형 및 비선형 layer가 포함된 L개의 layers로 구성된 네트워크가 있다고 가정해보자. 

 

학습데이터 D가 있을 때, 제약조건 C에서의 네트워크 가지치기 문제는 다음과 같이 공식화될 수 있다.

W^ : 가지치기 후 남은 파라미터 (W^ ⊂ W)

L : loss

f(·) : 네트워크 함수를 인코딩하는 함수

Φ(·) : 네트워크를 제약조건 C(ex. latency, FLOPs, memory)에 매핑하는 함수

->  L을 최소로 하는 W^를 찾는것, 이때, Φ(  f(·) )가 C를 만족해야함.

주로 이 task에서는 latency에 먼저 초점을 맞추고 다른 제약 조건을 고려한다.

        input_var = Variable(input)

        target_var = Variable(target)

        # 그냥 resnet 결과가 output인듯

        output = model(input_var)

        loss = criterion(output, target_var)

 

즉, 가장 중요한 것은 성능 하락을 최소한으로 하며, 제약조건을 충족하는 네트워크를 파악하는 것이다.

pl : l번째 layer에서 유지되는(=가지치기에서 살아남은) 뉴런의 수

Il(pl) : pl 뉴런을 사용했을 때의 최종 정확도에 대한 최대 중요성

Tl(pl−1, pl) : pl−1 input channels과 pl output channels을 사용하여 l번째 layer의 latency contribution

p0 ; 첫번째 conv block에 대한 고정된 input channel number (ex. 3 for RGB images)

 

 

Importance score

식 (2)에서의 Il(pl)를 개별 뉴런 Ppl j=1 I j l .으로부터의 누적 점수로 생각한다. loss 변화의 Taylor expansion을 사용하여 뉴런의 중요성을 근사한다 ( Importance estimation for neural network pruning ). 구체적으로 batch normalization layers를 잘라내고 I번째 layer에서 n번째 뉴런의 중요성은 다음과 같이 계산된다.

g : gradient of the weight

γ n l : batch normalization layer의 weight

β n l : batch normalization layer의 bias

( Importance estimation for neural network pruning )의 squared loss와 달리, 절대값의 차이를 이용하여 약간의 개선을 확인했다.

 

total importance를 최대화하기 위해, 가장 중요한 뉴런의 우선순위를 높게 설정한다. 이를 위해, l번째 layer에 있는 뉴런들을 중요도 점수에 따라 내림차순으로 순위를 매기고, j번째 순위에 있는 뉴런의 중요도 점수를 Ijl로 나타낸다.

            # Calculate the importance on batch normalization weights

            tmp_importance = 0

            for layer_name in group:

                layer = layers[layer_bn[layer_name]]

                tmp_importance += (

                    layer.weight.data * layer.weight.grad.data

                    + layer.bias.data * layer.bias.grad.data

                )

            importance_score = tmp_importance.abs()

 

 

 

*코드 에서의  loss는 다음과 같다. (어디랑 매칭되는건지 아직 못찾음ㅠㅠ)

    ##### Loss #####

    # 기본 loss는 CrossEntropyLoss

    loss = nn.CrossEntropyLoss

    # configs\exp_configs\rn50_imagenet_baseline.yaml 에서 mixup = 0.0

    if exp_cfg.mixup > 0.0:

        loss = lambda: NLLMultiLabelSmooth(exp_cfg.label_smoothing)

    # configs\exp_configs\rn50_imagenet_baseline.yaml 에서 label_smoothing = 0.1

    elif exp_cfg.label_smoothing > 0.0:

        loss = lambda: LabelSmoothing(exp_cfg.label_smoothing)

    criterion = loss()

    if cuda:

        criterion = criterion.cuda()

 

Latency contribution

미리 측정된 값을 활용하여 latency를 얻는다. 이 layer latency는 해당 layer에서의 각 뉴런의 latency contribution의 집합이다.

l번째 layer에서의 j번째 뉴런의 latency contribution은 아래 수식을 통해 계산할 수 있다.

                    latency = self._get_layer_latency(

                        cin=pre_active_neuron_num

                        if conv_groups == 1

                        else pre_active_neuron_num - channel_group_size * i,

                        cout=active_neuron_num - channel_group_size * i,

                        kernel=k,

                        feat_size=fmap,

                        stride=stride,

                        groups=conv_groups

                        if conv_groups == 1

                        else conv_groups - channel_group_size * i,

                    )

                    reduced_latency = self._get_layer_latency(

                        cin=pre_active_neuron_num

                        if conv_groups == 1

                        else max(pre_active_neuron_num - channel_group_size * (i + 1), 0),

                        cout=max(active_neuron_num - channel_group_size * (i + 1), 0),

                        kernel=k,

                        feat_size=fmap,

                        stride=stride,

                        groups=conv_groups

                        if conv_groups == 1

                        else max(conv_groups - channel_group_size * (i + 1), 0),

                    )

                    layer_latency_change = latency - reduced_latency

실제로 먼저, importance에 따라 뉴런의 순위를 매긴 후, latency contribution을 고려한다. 이는 다음과 같은 특성을 야기한다. 만약 layer l에서 가장 중요하지 않은 뉴런을 1개 삭제한다면, 뉴런의 수는 pl에서 pl-1로 변경되고, 이는 해당 뉴런의 latency contribution 점수가 c pl l 감소된다. 즉, 뉴런의 중요도에 따라 잠재적인 latency 가 감소된다. 해당 layer에서의 가장 높은 중요도를 갖는 뉴런은 항상 latency contribution c1l을 갖는다. latency와 accuracy를 고려한 reward maximization을 통해 이를 해결한다.

 

Augmented Knapsack Solver

importance와 latency 추정치가 구해진 후에, 식 (2)를 구하는 것을 목표로 한다. 

이를 통해 전체 가지치기 과정이 추가적인 선행 제약 조건만으로 knapsack 문제로 간소화된다. 선행 제약 조건은 l번째 layer에 순위가 j인 뉴런의 경우, 순위가 j보다 높은 뉴런이 l번째 layer에 유지되고 순위가 j보다 낮은 뉴런이 제거될 때만 해당 뉴런의 latency contribution이 유지된다는 사실에서 비롯된다.

 

하지만, 문제는 각 항목을 포함하기 전에 선택해야 하는 이전 항목 목록과 함께 지정하면 해결할 수 있습니다.

저희는 모든 선행 뉴런이 그 이전에 처리되도록 내림차순 중요도 점수로 정렬된 뉴런을 고려하기 위해 배낭 해결사를 증강합니다. 증강 배낭 해결사의 의사 코드에 대한 설명은 Algo. 1에 나와 있습니다(자세한 설명은 부록 A에 나와 있습니다). 증강 해결사는 지연 비용이 정확한지 확인해야 합니다.

 

 

 

 

코드 상에서의 학습 과정

main.py의 main()함수 실행

if __name__ == '__main__':

    torch.backends.cudnn.benchmark = True

    main()

 

main()에서 train_loop()함수 실행

    ##### Train #####

    best_prec1, best_prec5 = train_loop(

        exp_cfg,

        model,

        criterion,

        optimizer,

        lr_scheduler,

        train_loader,

        val_loader,

        batch_size_multiplier=batch_size_multiplier,

        train_writer=train_writer,

        pruner=pruner,

    )

 

training.py의 train_loop()에서 get_train_step()함수 실행

  * get_train_step은 함수 내 함수가 존재하는 코드임. 

    for epoch_id in range(exp_cfg.epochs):

        # Train one epoch

        losses = AverageMeter()

        top1 = AverageMeter()

        top5 = AverageMeter()

        data_time = AverageMeter()

        compute_time = AverageMeter()

        ips = AverageMeter()

        model.train()

        step = get_train_step(

            model, criterion, optimizer, exp_cfg.fp16, use_amp=exp_cfg.amp, pruner=pruner

        )

 

get_train_step()에서 _step 함수 객체를 바인딩함

def get_train_step(model, criterion, optimizer, fp16, use_amp, pruner=None):

    def _step(input, target, global_step, optimizer_step=True):

                     ...

    return _step

 

step 객체를 통해 get_train_step()내부의 _step 함수 실행

        for i, (input, target) in enumerate(train_loader):

            model.train()

            lr_scheduler(optimizer, epoch_id)

            data_load_time = time.time() - end

            optimizer_step = ((i + 1) % batch_size_multiplier) == 0

            loss, prec1, prec5 = step(input, target, global_iter, optimizer_step=optimizer_step)

 

model 추론 후 loss 계산

    def _step(input, target, global_step, optimizer_step=True):

        optimizer.zero_grad()

        input_var = Variable(input)

        target_var = Variable(target)

        output = model(input_var)

        loss = criterion(output, target_var)

 

loss.backward()로 gradient 계산

        if fp16:

            optimizer.backward(loss)

        elif use_amp:

            with amp.scale_loss(loss, optimizer) as scaled_loss:

                scaled_loss.backward()

        else:

            loss.backward()

 

pruning을 여기서 하는거같다?

        if pruner is not None:

            pruner.update_metric(global_step)

 

optimizer.step()으로 파라미터 업데이트

        if optimizer_step:

            optimizer.step()

 

 

코드 상에서의 Pruning 부분

main.py의 main()함수 실행 // 학습 과정에 pruning이 존재함. no_prune이면,,??

    ##### Pruner #####

    if args.no_prune:

        pruner = None

    else:

        pruner = Pruner(model, exp_cfg)

 

prune_step -> _latency_target_pruning -> get_total_latency

prune_step -> get_average_importance -> 

    def prune_step(self, global_step):

        pruned_num = 0

        if (

            self._prune_start <= global_step < self._prune_end and

            (global_step + 1 - self._prune_start) % self._prune_interval == 0

        ):

            logger.info("Performing pruning.")

            self._importance_calculator.get_average_importance()

            pruned_num = self._latency_target_pruning(target_latency=self._prune_target.pop(0))

            self.save_mask(os.path.join(os.environ["OUTPUT_PATH"], "group_mask.pkl"))

            self._importance_calculator.reset_importance()

        else:

            self.mask_weights()

       

        return pruned_num