PyTorch 파이닝 — torch.nn.utils.prune으로 희소화하기
PyTorch 파이닝 — torch.nn.utils.prune으로 희소화하기
PyTorch 튜토리얼은 torch.nn.utils.prune 로 신경망을 희소화하는 법을 알려줘요. 특히 완전 연결·합성곱 계층의 불필요한 연결을 제거해 경량 모델을 만들어 기기(디바이스)에 배포할 수 있게 해요.
어떻게 가지치나
모듈과 파라미터 이름, 가지치기 기법, 비율(amount)을 지정해요.
import torch.nn.utils.prune as prune
prune.random_unstructured(module.conv1, name="weight", amount=0.3)
amount 가 0~1 사이면 연결의 퍼센트, 정수면 절대 개수를 의미해요. 희소 부분망(lottery tickets) 연구 등에도 활용돼요.
재매개변수화 방식
가지치기는 원래 파라미터를 weight_orig 로 보존하고 마스크(weight_mask)를 곱해 재매개변수화해요. 때문에 순수한 크기 감소는 아니고, 연구·반복 가지치기에 적합해요. 배포 시에는 prune.remove() 로 마스크를 원 파라미터에 확정(fold)하는 걸 권장해요.
커스텀 기법 확장
BasePruningMethod 를 상속해 나만의 기법(compute_mask/apply 구현)을 만들 수도 있어요.