Project/Proj-Nova Vision

Nova-Vision은 어떻게 개발했을까?: 모델 경량화 및 코드 최적화

dohyeon2 2024. 11. 8. 19:41

목차

    이번 글에서는 "Nova Vision - 농장주를 위한 가축 성장관리 서비스"의 고도화 과정에서 마주했던 문제들과 해결방법에 대해 정리하였습니다. Nova Vision 서비스의 MVP를 완성하고 처음 테스트 했을 때 다음과 같은 문제점들을 마주하였습니다.

    • 이미지 한장 당 Inference time이 10초 이상 소요
    • AWS 온디맨드 요금이 요청 한번에 0.035$ 발생

    이러한 문제들을 해결하고자 코드 최적화 -> 연산 최적화 -> 모델 경량화 순서로 최적화를 진행하였습니다.

     

    개발과정
    [1] Nova-Vision은 어떻게 개발했을까?: 3D Depth Camera 개발

    [2] Nova-Vision은 어떻게 개발했을까?: AI 가축 체중측정 알고리즘

    [3] Nova-Vision은 어떻게 개발했을까?: 모델 경량화 및 코드 최적화

    Github

    [1] https://github.com/dohyeonYoon/Nova-Vision

    [2] https://github.com/dohyeonYoon/BRODY

     

     

    코드 최적화


    모델 경량화나 연산 최적화는 모델 자체를 tuning하여 전체적인 inference time을 줄이는 방법입니다. 하지만 모델을 아무리 압축하고 연산을 최적화하더라도 코드가 최적화 되어있지 않다면 전체 프로그램의 실행속도를 줄이는데 한계가 있고, 리소스 사용량이 늘어날 수밖에 없습니다. 또한 코드 최적화는 단순히 프로그램의 실행속도를 빠르게 하는 것 외에도, 코드의 가독성과 유지보수성도 높여주기 때문에 이후 모델 경량화나 연산 최적화 과정에서 디버깅이 더 용이할 것이기 때문에 코드 최적화를 가장 먼저 진행하였습니다.

     

    1. 병목지점 찾기

    전체 실행시간 측정하기

    코드 프로파일링에 앞서 python time 라이브러리를 활용하여 전체 코드시간을 측정한 결과, 한 건의 데이터를 처리하는데 약 10.1023초가 소요되었습니다.

    python pyBRODY.py

     

    라인별 실행시간 측정하기(디테일한 프로파일링)

    line_profiler 라이브러리를 사용하여 line by line으로 실행시간을 측정하였습니다.

    • 총 실행시간: 10.2269초(프로파일링 오버헤드 시간 포함)
    • 첫번째 병목지점: 실행시간-4.2393초, 전체 비율-41.5%, 위치-line 35
    • 두번째 병목지점: 실행시간-0.3810초, 전체 비율-3.7%, 위치-line 36
    • 세번째 병목지점: 실행시간-0.3704초, 전체 비율-3.6%, 위치-line 40
    • 네번째 병목지점: 실행시간-1.8711초, 전체 비율-18.3%, 위치-line 41
    • 다섯번째 병목지점: 실행시간-0.2940초, 전체 비율-2.9%, 위치-line 44
    • 여섯번째 병목지점: 실행시간-2.6720초, 전체 비율-26.1%, 위치-line 56

    위 여섯개의 병목지점에서 전체 96.1%의 실행시간을 잡아먹었습니다. 

     

    2. 병목지점 해결하기

    저는 병목지점을 해결하기 위해 다음과 같은 사항들을 고려하였습니다.

    1. 불필요한 계산 없애기: 기본적으로 반복되거나 불필요한 계산은 프로그램의 실행속도를 늦추고, 메모리 사용량을 증가시키기 때문에 이 부분을 중점적으로 살펴보았습니다.
    2. 루프 최적화하기: 지금까지 경험상 for문과 같은 반복루프에서 가장 많은 시간을 차지하기 때문에 이 부분을 최적화 하면 실행시간을 줄일 수 있다고 생각했습니다.
    3. 적절한 자료구조와 알고리즘 선택하기: 저처럼 3D array를 빈번하게 인덱싱할 경우 메모리 측면에서 매우 불리합니다. Flattening, Octree등의 자료구조를 사용하여 3D 데이터 탐색시간을 줄일 수 있다고 생각했습니다. 

    첫번째 병목지점(Seg.Segment_Broiler 함수, line 35) 해결하기

    Seg.Segment_Broiler 함수 line by line 실행시간

    • line 29(모델 불러오기): 코드 최적화의 영역이 아닌 모델 경량화와 연산 최적화의 영역이므로 제외
    • line 35(모델 추론하기): 코드 최적화의 영역이 아닌 모델 경량화와 연산 최적화의 영역이므로 제외
    • line 32(이미지 읽기): cv2.imread() 함수를 사용하여서 Pillow Image.open() 함수로 대체하려고 했으나, mmdetection inference 함수의 input format은 numpy array이어야 하고, BGR 순서여야 했습니다. 따라서 cv2.imread()와 Pillow Image.open()+np.array()+RGB2BGR의 처리속도를 비교하였습니다.
      소요시간(초) 증감률
    cv2.imread() 0.3474 0%
    pillow.Image.open()
    +np.array()
    +RGB2BGR
    0.4020 13.6% 증가

    결과적으로 기존 opencv imread 함수가 이미 inference 함수에 최적화 되어있기 때문에 기존 코드를 유지하기로 결정하였습니다.

     

    두번째 병목지점(Seg.Get_mask 함수, line 36) 해결하기

    Seg.Get_mask 함수 line by line 실행시간

    • line 57: line58에서 cv2.findNonZero 메소드로 이미 0이 아닌 픽셀을 찾고있으므로 line 57은 불필요한 line인 것을 확인하였습니다. 따라서 line 57을 삭제하고, line58에 results[1][0][i].astype(np.uint8)로 대체하였습니다.
      개선 전 개선 후 증감률
    line 57삭제 0.3873 초 0.1234 초 0.2639초 감소
    (68.2% 개선)

     

    세번째 병목지점(Conv.Generate_Depthmap_1 함수, line 40) 해결하기

    Conv.Generate_Depthmap_1 함수 line by line 실행시간

    • line 104~6: 리스트 컴프리헨션을 이용하여 이중 for문과 append 메소드를 한줄로 변환하고 최적화하였습니다.
      개선 전 개선 후 증감률
    line 104~6 0.5455 초 0.2519 초 0.2936초 감소
    (54% 개선)

     

     

    네번째 병목지점(Conv.Convert_3D, line 41) 해결하기

    Conv.Convert_3D 함수 line by line 실행시간

    • line 146~154: 기존 이중루프로 width x height개의 모든 픽셀을 하나씩 탐색하며 X,Y,Z값을 계산하던 코드를 numpy.meshgrid()를 활용하여 모든 픽셀을 한번에 계산하는 방식으로 개선하였습니다.

    • line 160~7: 리스트 컴프리헨션을 사용하여 for문을 한줄로 만들어서 연산과정을 최적화했습니다.

      개선 전 개선 후 증감률
    line 146~154 2.3097 초 0.2088 초 2.1초 감소
    (91% 개선)
    line 160~7 0.2088 초 0.1431 초 0.0657초 감소
    (31.5% 개선)

     

    다섯번째 병목지점(Exclu.Remove_outlier 함수, line 44) 해결하기

    Exclu.Remove_outlier 함수 line by line 실행시간

    • line 47:  추론에 불필요한 코드 제거
      개선 전 개선 후 증감률
    line 146~154 0.2975 초 0.2327 초 0.0648초 감소
    (22% 개선)

     

    여섯번째 병목지점(Visual.Build_PNG 함수, line 56) 해결하기

    개선 전 - Build_PNG 함수 line by line 실행시간

    • line 64(모델 불러오기): 앞서 Segment_Broiler 함수에서 한번 model을 불러왔는데 중복으로 한번 더 불러오는 것을 발견하였습니다. 그래서 저는 Segment_Broiler 함수에서 model을 반환하고 Build_PNG에서 model을 인자로 받아오게 함으로써 한번만 호출하도록 코드를 수정하였습니다. 또한 이에따라 mmdet 라이브러리 import하는 line을 제거하였습니다.
    • line 82(추론결과 시각화하기): model.show_result 함수를 뜯어봤을 때, imshow_det_bboxes 함수가 대부분의 실행시간을 차지하는 것을 확인하였습니다. 그중 line 402의 print_to_buffer 메소드가 약 0.9초로 가장 많은 시간을 소비하였는데요. 이는 matplotlib 라이브러리 매소드입니다. 해당 부분을 최적화하기 위해 저는 matplotlib.use('Agg') 매서드를 사용하여 저처럼 GUI가 필요하지 않은 서버환경에서 matplotlib을 백엔드로 실행하여 실행시간을 최적화 하였습니다.

      개선 전 개선 후 증감률
    model 2번호출 -> 1번 호출 10.2269 초 9.3327 초 0.8942초 감소
    (8.8% 개선)
    matplotlib 백엔드 실행 1.7595 초 1.7435 초 0.016초 감소
    (1% 개선)

     

    위와 같이 총 6개의 병목지점을 해결하고 추론시간을 총 2.87초 개선하였습니다.
    (기존 10.1023초에서 7.22402초로)

     

    Reference

    [1] https://opencvpython.blogspot.com/2012/06/fast-array-manipulation-in-numpy.html

    연산 최적화(with TensorRT)


    이번에는 PyTorch로 학습된 Instance Segmentation 모델을 TensorRT 엔진으로 변환하는 과정을 살펴보겠습니다.

    저는 PyTorch의 Mask-RCNN + ResNet-101 모델을 사용하였으며 [표 1]과 같은 환경에서 테스트하였습니다.

    표 1. 테스트 환경

     

    PyTorch 모델을 TensorRT로 추론하는 두가지 방법

    2024년 10월 23일 기준 PyTorch 모델을 직접 TensorRT 엔진으로 변환하는 도구는 없습니다. 대신 PyTorch 모델을 TensorRT Runtime 으로 추론할 수 있는 두 가지 방법이 있습니다. 첫 번째는 ONNX 모델로 변환 후 TensorRT 엔진으로 변환하는 방법이고, 두 번째는 TorchScript로 변환 후 Torch-TensorRT를 이용하는 방법입니다. 이 두 방법의 특징은 다음과 같습니다.

     

    1. ONNX 모델 변환 후 TensorRT 엔진 변환

    이 방법은 PyTorch로 학습된 모델을 TensorRT 엔진으로 변환하는 가장 일반적인 방법입니다. ONNX는 여러 ML 프레임워크에서 만들어진 모델들을 서로 호환되도록 만들어주는 일종의 공유 플랫폼입니다. PyTorch도 ONNX로 변환이 가능하고, TensorRT는 다시 ONNX로부터 변환하여 사용할 수 있습니다. 이 방법의 장점은 최적화가 잘 된 온전한 TensorRT 엔진을 사용할 수 있다는 것입니다. 그러나 PyTorch 모델 내에 ONNX나 TensorRT가 지원하지 않는 연산자가 포함된 경우에는 모델 수정이 필요하다는 단점이 있습니다.

     

    2. TorchScript로 변환 후 Torch-TensorRT를 이용

    Torch-TensorRT는 TorchScript 모델을 TnesorRT를 이용해 최적화할 수 있도록 해주는 일종의 컴파일러입니다. TorchScript는 PyTorch 모델의 파이썬에 대한 의존성을 없애고, 다양한 환경에서 고성능 추론이 가능하도록 모델을 컴파일한 포맷입니다. 이러한 TorchScript 모델을 Torch-TensorRT를 이용해 변환하면 TensorRT와 호환되는 연산자는 TensorRT 연산자로, 호환되지 않는 연산자는 PyTorch 연산자로 처리합니다. 따라서 이 방법은 ONNX를 사용한 방법보다 호환성 문제가 적다는 장점이 있습니다. 그러나 [그림 1]과 같이 TorchScript 모델 내에서 PyTorch 연산자와 TensorRT 연산자가 함께 사용되기 때문에 온전한 TensorRT 엔진에 비해서 추론속도가 느리며, TensorRT와 호환되는 연산자의 갯수가 적을수록 속도는 더 느려집니다. 

    그림 1. Torch-TensorRT의 런타임 구조

    따라서 저처럼 추론속도를 최대한 줄여야하는 상황에서는 온전한 TensorRT 엔진을 사용하는 "ONNX 모델 변환 후 TensorRT 엔진 변환" 방식을 사용하시는 것을 추천드립니다. 이어서 PyTorch->ONNX->TensorRT로 변환하는 자세한 과정을 설명드리겠습니다.

     

    PyTorch 모델을 ONNX 모델로 변환

    import torch
    
    # Load my own Instance Segmentation model
    check_file = './mmdetection/weights/mask_rcnn_r101_3x.pth'
    model = torch.load(check_file)
    model.eval()
    
    # Dummy input with the model input size.
    dummy_input = torch.randn(1, 3, 720, 1280)
    
    # Save ONNX model to ONNX_PATH.
    torch.onnx.export(model, 
                      dummy_input, 
                      ONNX_PATH, 
                      opset_version=13,
                      input_names=["input"], 
                      output_names=["dets", "labels", "masks"])

     

    ONNX 모델을 TensorRT 엔진으로 변환

    ONNX 모델을 TensorRT 엔진으로 변환하기 위해 저는 Polyhraphy 라이브러리를 사용했습니다. Polygraphy는 NVIDIA에서 제공하는 개발 도구인데요. TensorRT API를 기반으로 만들어진 고수준(high-level) 추상화 인터페이스입니다. 따라서 TensorRT API보다 사용이 간편하고, 버전에 따른 호환성 문제도 더 적습니다(안쓸 이유가 없죠?). 또한 TensorRT API에 기반하여 개발되었으므로 Polygraphy와 TensorRT API를 혼용하여 사용할 수 있다는 장점도 갖고있습니다.

     

    이제 Polygraphy 라이브러리를 이용하여 ONNX 모델을 TensorRT 엔진으로 변환하는 방법을 알아보겠습니다.

     

    1. Calibrator 객체 생성

    제가 진행하는 모델 경량화 기법은 model quantization, 정확하게는 훈련 후 양자화 기법입니다. 이를 수행하려면 Polygraphy의 Calibrator  클래스를 사용해야 합니다. Calibrator 클래스는 입력 데이터를 generator의 형태로 받습니다. 따라서 추론 환경과 동일한 전처리가 적용된 입력 데이터를 생성하는 generator를 만들고, 이를 Calibrator에게 인자로 전달해야 합니다.

    import os
    
    from PIL import Image
    from polygraphy.backend import trt as poly_trt
    
    # Polygraphy needs a generator-type data loader.
    val_list = os.listdir(IMAGE_DIR)
    def data_generator():
        for image in val_list:
            # Preprocess.
            image_path = os.path.join(IMAGE_DIR, image)
            image = Image.open(image_path).convert("RGB")
            # Add batch dimension.
            image = image.unsqueeze(0)
            # Polygraphy uses numpy input.
            image = image.numpy()
            # Dict key must be the same as ONNX input name.
            yield {"input": image}
    
    calibrator = poly_trt.Calibrator(data_loader=data_generator())

    이때 data_generator는 IMAGE_DIR에서 입력 이미지를 가져와서 앞서 만든 ONNX 모델의 입력 형식과 동일한 형식의 dict를 반환합니다. 이후 Calibrator 타입의 객체를 만드는데, 이때 위에서 만든 generator를 인자로 사용합니다.

     

    2. ONNX 모델 로드 및 IBuilderConfig 객체 생성

    Clibrator를 생성한 후, ONNX 모델을 로드하고, TensorRT 엔진 빌드에 필요한 각종 설정정보를 담고있는 IBuilderConfig 객체를 만들어야 합니다.

    builder, network, parser = poly_trt.network_from_onnx_path(path=BRODY/src/mmdetection/weights/mask_rcnn_r101_3x.onnx)
    
    # Builder config 생성
    builder_config = poly_trt.create_config(builder=builder,
                                            network=network,
                                            int8=True,
                                            fp16=True,
                                            calibrator=calibrator)

    Polygraphy network_from_onnx_path를 이용하여 ONNX 모델을 로드하면 builder, network, parser 세가지 객체가 반환됩니다. 이 중 builder network IBuilderConfig 객체를 만드는 create_config의 인자로 입력됩니다.

    이후 앞서 만든 Calibrator 객체를 인자로 추가해주면 IBuilderConfig 객체 생성이 완료됩니다.

     

    # 주의사항

    IBuilderConfig 객체를 만들 때에는 정수 타입을 지원하지 않는 레이어를 위해 FP16 타입 변환에 대한 옵션을 True로 설정해줘야 합니다.

     

    3. TensorRT 엔진 빌드 및 저장

    Polygraphy에서 TensorRT 엔진을 빌드하고 저장하는 방법은 다음과 같습니다.

    engine = poly_trt.engine_from_network(network=(builder, network, parser),
                                          config=builder_config)
    
    # TensorRT engine 저장
    poly_trt.save_engine(engine, "BRODY/src/mmdetection/weights/mask_rcnn_r101_3x.trt")

    앞서 ONNX 모델에서 로드된 builder, network, parser 객체들과 IBuilderConfig 객체를 engine_from_network에 인자로 입력하면, TensorRT 엔진이 빌드됩니다.

    마지막으로 빌드된 엔진 객체를 save_engine 함수를 이용하여 저장하면 ONNX to TensorRT 엔진 변환이 완료됩니다!

     

    4. TensorRT 엔진 추론

    저장된 TensorRT 엔진을 로드하여 본인의 데이터 추론을 진행하는 방법은 다음과 같습니다.

    # Load serialized engine using 'open'.
    engine = poly_trt.engine_from_bytes(open(ENGINE_PATH, "rb").read())
    
    with poly_trt.TrtRunner(engine) as runner:
        # Preprocess.
        image = transform(Image.open(IMAGE_PATH).convert("RGB"))
        image = image.unsqueeze(0).numpy()
        # Input dict keys are the same as 'input_names' arg in 'torch.onnx.export'.
        output_dict = runner.infer({"input": image})
        # Output dict keys are the same as 'output_names' arg in 'torch.onnx.export'.
        output = output_dict["output"]

     

    engine_from_bytes를 이용하여 TensorRT 엔진을 로드하고, TrtRunner를 사용하여 runner 객체를 생성합니다. 이 runner 객체는 추론 시 runner.infer 를 사용하고, 입력과 출력 모두 dict 타입을 사용합니다. 이때 입력 dict의 key로는 ONNX 모델 생성 시 torch.onnx.export에 인자 input_names로 전달했던 값을, 출력 dict의 key로는 인자 output_names로 전달했던 값을 사용해야만 정상적으로 추론이 진행됩니다!! 즉, 이미지를 전처리한 후 {input_names:image} 형태의 dict로 변형하여 runner.infer에 입력해줘야 합니다. 또한 출력 dict의 key로 output_names를 이용해 value에 접근하여 모델의 출력을 얻을 수 있습니다.

     

     

    모델 경량화


    모델 경량화에는 크게 세 가지 방법이 존재합니다. 각각 Pruning, Quantization, Distillation인데요. 저는 이 세 가지 방법을 제 Instance Segmentation 모델에 적용한 과정을 소개해보겠습니다.

     

    작성중...