[논문 리뷰] Streamlining Tensor and Network Pruning in PyTorch

M. Paganini, Jessica Zosa Forde|arXiv (Cornell University)|2020. 04. 28.

Computational Physics and Python Applications참고 문헌 11인용 수 7

한 줄 요약

이 논문은 훈련, 추론 또는 훈련 후 단계에서 신경망 레이어에 구조적 및 비구조적 프루닝을 적용하기 위한 통합된, 오픈소스 인터페이스인 PyTorch `torch.nn.utils.prune` 모듈을 소개한다. 이는 연구자와 실무자들이 최소한의 코드 변경으로 모델 크기와 계산량을 줄일 수 있도록 하며, 반복적 프루닝, 전역 크기 비교, 프루닝된 모델의 간편한 직렬화를 지원하는 일관된 API를 제공한다.

ABSTRACT

In order to contrast the explosion in size of state-of-the-art machine learning models that can be attributed to the empirical advantages of over-parametrization, and due to the necessity of deploying fast, sustainable, and private on-device models on resource-constrained devices, the community has focused on techniques such as pruning, quantization, and distillation as central strategies for model compression. Towards the goal of facilitating the adoption of a common interface for neural network pruning in PyTorch, this contribution describes the recent addition of the PyTorch torch.nn.utils.prune module, which provides shared, open source pruning functionalities to lower the technical implementation barrier to reducing model size and capacity before, during, and/or after training. We present the module's user interface, elucidate implementation details, illustrate example usage, and suggest ways to extend the contributed functionalities to new pruning methods.

연구 동기 및 목표

• 모바일, IoT, AR/VR 시스템과 같은 자원이 제한된 장치에 대규모 과도하게 파rameter화된 딥러닝 모델을 구현하는 데 점점 커지는 과제를 해결하기 위해.
• PyTorch 내부에서 공통의 오픈소스 인터페이스를 제공함으로써 모델 프루닝을 구현하는 데 기술적 장벽을 낮추기 위해.
• 연구자들이 공통의 API를 통해 새로운 프루닝 기법을 쉽게 실험하고 기여할 수 있도록 하기 위해.
• 일관되고 모듈적이며 확장 가능한 설계 원칙을 바탕으로 훈련 중 및 훈련 후 프루닝을 모두 지원하기 위해.
• 장치 내에서의 추론을 통해 효율성 향상, 에너지 소비 감소 및 프라이버시 향상을 위한 모델 압축을 촉진하기 위해.

제안 방법

• 모든 프루닝 기법에 대한 공통 인터페이스를 정의하는 추상 기반 클래스인 `BasePruningMethod` 를 도입하며, `compute_mask` 의 구현을 요구한다.
• 재구성 기반 기법: 프루닝된 파라미터를 원본 텐서를 `name_orig` 으로, 마스크를 `name_mask` 로 모듈 버퍼에 저장함으로써 마스크된 버전으로 대체한다.
• 정방향 프리훅을 사용하여 정방향 전파 중에 원본 텐서를 마스크로 동적으로 곱함으로써 계산 그래프의 무결성을 유지한다.
• 구조적 및 비구조적 프루닝을 모두 지원하며, `L1Unstructured`, `RandomUnstructured`, `LnStructured` 와 같은 전용 클래스를 통해 설정 가능한 프루닝 비율과 차원을 제공한다.
• `PruningContainer` 를 사용하여 동일한 파라미터에 대해 여러 번의 프루닝 작업을 추적함으로써 반복적 프루닝을 가능하게 한다.
• `prune.global_unstructured` 와 같은 유틸리티 함수를 제공하여 전체 네트워크에 걸쳐 전역 크기 기반 프루닝을 수행하며, 모든 파라미터를 취합하여 비교한다.

실험 결과

연구 질문

• RQ1 PyTorch 내부에서 다양한 프루닝 전략을 지원할 수 있는 통합적이고 확장 가능하며 사용자 우량한 프루닝 인터페이스를 설계하는 방법은 무엇인가?
• RQ2 PyTorch의 autograd 및 직렬화 워크플로우에 원활하게 통합될 수 있도록 안전하고 되돌릴 수 있으며 조합 가능한 프루닝 작업을 가능하게 하는 아키텍처 패턴은 무엇인가?
• RQ3 전체 모델에 걸친 전역 프루닝을 효율적으로 구현하면서도 레이어 단위 및 반복적 프루닝과의 호환성을 유지하는 방법은 무엇인가?
• RQ4 프루닝된 모델이 손실 없이 영구적으로 저장되거나 복원될 수 있도록 보장하는 메커니즘은 무엇인가?
• RQ5 연구자들이 PyTorch의 모듈 시스템 내부 지식 없이도 쉽게 새로운 프루닝 방법을 구현하고 기여할 수 있도록 API를 어떻게 구성할 수 있는가?

주요 결과

• `torch.nn.utils.prune` 모듈은 최소한의 코드 변경으로 PyTorch에서 구조적 및 비구조적 프루닝을 일관되고 오픈소스 방식으로 적용할 수 있도록 한다.
• 프루닝 작업은 PyTorch의 autograd 시스템과 완전히 호환되며, 훈련 이전, 동안, 이후 어느 시점에서나 적용 가능하며 결과는 모델의 `state_dict` 에 그대로 유지된다.
• 모듈은 `global_unstructured` 를 통해 전체 네트워크에 걸쳐 전역 프루닝을 지원하며, 모든 레이어에 걸쳐 연결 수의 하위 20%를 크기 기반으로 프루닝할 수 있다.
• 반복적 프루닝은 `PruningContainer` 를 사용하여 동일한 파라미터에 대해 반복적으로 적용 가능하며, 예를 들어 3개의 항목을 프루닝한 후 나머지 채널의 50%를 프루닝하는 점진적 압축 전략을 구현할 수 있다.
• 모듈은 하드 프루닝(이진 마스크)과 소프트 프루닝을 모두 지원하며, `prune.remove` 를 통해 영구적으로 프루닝을 제거하여 프루닝된 텐서를 원래 파라미터 이름으로 복원할 수 있다.
• 설계는 프루닝된 모델의 직렬화 및 역직렬화를 원활하게 가능하게 하여 표준 PyTorch 모델 저장 및 로딩 워크플로우와의 호환성을 보장한다.