[최적화] 학습 속도 개선(AMP, Prefetch)

서론

딥러닝 모델의 학습/추론 프로세스에는 많은 단계가 포함됩니다. 각 실험의 반복이 빠를수록 제한된 시간과 리소스로 더 효율적인 성능 최적화를 할 수 있습니다. 아래 과정으로 성능 최적화 기법을 적용해 보겠습니다.

 

본론

- 모델 설정

backbone : vit-base_path16_224
input size : [224, 224]
epoch : 1

무거운 모델이 성능 면에서 비교하기 쉬울 것 같아서 VIT로 선택했습니다. 1 에폭만 돌려서 학습 속도의 차이를 확인해 보겠습니다.

 

- 기본 Baseline 코드

기법을 적용하지 않은 학습 결과입니다. 1 에폭이 동작한 시간과 사용한 메모리 양을 나타내었습니다.

 

- With AMP

-> AMP란?

{"originWidth":485,"originHeight":145,"style":"alignCenter","caption":"[ \n bfloat16 숫자 형식

모델의 파라미터를 32-bit가 아닌 16-bit로 표현하여 배치 사이즈를 늘리고, 그에 따라 학습 속도를 빠르게 할 수 있는 기술입니다.

 

1) 같은 모델이라도 더욱 빠른 학습을 통해 빠르게 결과를 얻을 수 있음
2) 빠른 학습, 그에 따른 적은 GPU 사용량이 이산화탄소 배출 저감으로 이어짐
3) 모델 구조에 구애받지 않고 모든 모델에 적용될 수 있음

 

-> Mixed Precision Training

1) FP32로 표현된 FP32 Master weights을 복사하여 FP16 weights를 만든다.
2) FP16으로 Forward Propagation 진행 (gradient는 FP16일 것)
3) Loss 값을 Scale factor S로 곱한다
4) 얻은 FP16 gradient를 Backpropagate 한다
5) 3번에서 Loss를 S로 곱했으니 Backprop해서 얻은 weight gradient에 S를 나눈다
6) Gradient clipping, weight decay 등을 적용한다
7) FP32 Master weights을 업데이트한다

 

-> 실험

scaler = torch.cuda.amp.GradScaler()

for idx, train_batch in enumerate(train_loader):

    inputs, labels = train_batch
    inputs = inputs.to(device)
    labels = labels.to(device)

    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        outs = model(inputs)
        preds = torch.argmax(outs, dim=-1)
        loss = criterion(outs, labels)
	
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

파이토치는 학습 최적화 기법을 라이브러리로 제공합니다. base line 코드에서 scaler 부분을 추가해 주면 됩니다.

결과 : 속도 약 50% 감소, 메모리 사용량 감소, 인퍼런스 성능 감소

주의 깊게 볼 점은 AMP 등 성능 최적화 기법을 사용하면 인퍼런스 성능이 감소할 수 있다는 점입니다. 따라서 많은 실험을 통해 학습에만 AMP를 적용할지, 결정해야 합니다.

 

- with AMP + zero_grad

-> zero_grad란?

optimizer에서 gradients를 0이 아닌 None으로 적용하는 방법으로 None일 경우 else 밑에 연산을 하지 않아, 학습 시간을 줄일 수 있습니다.

 

-> 실험

scaler = torch.cuda.amp.GradScaler()

for idx, train_batch in enumerate(train_loader):

    inputs, labels = train_batch
    inputs = inputs.to(device)
    labels = labels.to(device)

    optimizer.zero_grad(set_to_none=True) # 여기
    with torch.cuda.amp.autocast():
        outs = model(inputs)
        preds = torch.argmax(outs, dim=-1)
        loss = criterion(outs, labels)
	
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

optimizer의 gradient를 None으로 지정하고 학습을 진행합니다.

결과 : 속도 감소, 메모리 사용량 감소

속도와 메모리 모두 감소하였으며, 추론 성능은 동일합니다.

 

- with AMP + zero_grad + non_blocking

-> non_blocking이란?

{"originWidth":589,"originHeight":465,"style":"alignCenter","caption":"[\n Optimize PyTorch Performance for Speed and Memory Efficiency (2022)

비동기적으로 데이터를 CPU에서 GPU로 넘겨주는 방식입니다. 데이터가 GPU로 올라올 때 순차적이 아닌 병렬적으로 데이터 전송을 할 수 있습니다. 따라서 GPU에서 데이터를 많이 다룬다면 꼭 사용하면 좋을 것입니다.

 

-> 실험

scaler = torch.cuda.amp.GradScaler()

for idx, train_batch in enumerate(train_loader):

    inputs, labels = train_batch
    inputs = inputs.to(device, non_blocking=True) # 여기
    labels = labels.to(device, non_blocking=True)

    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        outs = model(inputs)
        preds = torch.argmax(outs, dim=-1)
        loss = criterion(outs, labels)
	
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

데이터를 불러오는 과정에서 non_blocking을 True로 세팅하고 학습합니다.

결과 : 속도 감소, 메모리 사용량 증가

속도는 감소한 반면 메모리 사용량이 증가하였습니다. 기법을 적용한다고 해서 반드시 성능이 좋아지는 것은 아니며 따라서 많은 실험을 통해 트레이드오프를 체크해야 합니다.

 

- with AMP + zero_grad + non_blocking + Fast prefetcher

-> Fast prefetcher란?

pytorch에서는 cuda stream과 같이 cpu에서 gpu로 데이터의 흐름을 제어할 수 있는 기능을 제공합니다.

 

class PrefetchLoader():
    def __init__(self,
                 loader,
                 mean = (0.548, 0.504, 0.479),
                 std = (0.237, 0.247, 0.246)) -> None:
        self.loader = loader
        self.mean = torch.tensor([x * 255 for x in mean]).cuda().view(1, 3, 1, 1)
        self.std = torch.tensor([x * 255 for x in std]).cuda().view(1, 3, 1, 1)
    
    def __iter__(self):
        stream = torch.cuda.Stream()
        first = True
        
        for next_input, next_target in self.loader:
            with torch.cuda.stream(stream):
                next_input = next_input.cuda(non_blocking=True)
                next_input = next_input.float().sub_(self.mean).div_(self_std)
                next_target = next_target.cuda(non_blocking=True)

            if(not first):
                yield input, target
            else:
                first = False
                
            torch.cuda.current_stream().wait_stream(stream)
            input = next_input
            target = next_target
            
        yield input, target

생성자를 보면 정규화가 데이터 셋이 아닌 로더에서 적용되어 있습니다. 쿠다의 흐름 안에서 GPU 연산으로 정규화를 진행하면 좀 더 연산 속도가 빨라집니다. 데이터 증강은 CPU에서 일어나는데 GPU에서 하게 하면 더 빠른 학습을 할 수 있습니다. 이런 기법을 사용한다면 CPU에서 하던 연산을 GPU로 옮기는 것이 더 효율적입니다.

 

-> 실험

결과 : 속도 감소, 메모리 사용량 감소

 

결론

최종적으로 학습 시간과 메모리 사용량을 효과적으로 사용할 수 있습니다. LLM과 같이 점점 더 큰 모델이 나올 텐데 사소한 부분일지라도 반드시 알고 기본으로 가지고 가야 합니다.