[최적화] AMP 이슈 및 피클 직렬화로 속도 높이기

서론

대회로 수행하고 있는 EAST 모델은 데이터를 로드하는 과정에서 geometry map, scoremap을 만들기 때문에 적은 데이터임에도 불구하고 1에폭당 10분이 넘게 걸렸습니다. 이를 해결하는 과정을 기록합니다. 

 

본론

- 데이터 로딩 과정

EAST 모델 데이터 셋의 __getitem__은 위처럼 전처리된 데이터를 반환하기에 로딩 시간이 굉장히 오래 걸렸습니다.

 

- 최적화로 시간 줄이기

1. 기본 소요 시간

제공받은 베이스라인 코드는 1 에폭에 13분이 소요됩니다.

 

2. AMP 적용

with torch.cuda.amp.autocast():
    loss, extra_info = model.train_step(img, gt_score_map, gt_geo_map, roi_mask)
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

첫 번째로 적용한 기법은 AMP입니다. 학습 코드에 AMP를 적용해 주었습니다.

 

=> AMP 이슈

하지만, AMP를 적용하니 IOU loss가 NAN인 문제가 발생하였습니다.

 

-> IOU loss 디버깅 과정

현재 사용하고 있는 loss에서 iou_loss는 geometry loss를 구할 때 발생합니다. EAST는 각 화소가 글자 영역 중심에 속하는 지와 글자 영역이라면 해당 bbox의 위치는 어디인지를 구하여 학습하게 되는 데 그중 해당 화소가 bbox와 얼마나 겹치는지 구하는 geometry map loss에서 얼마나 겹치는 지 구하는 iou loss를 반환하게 됩니다.

 

iou_loss_map = -torch.log((area_intersect + 1.0) / (area_union + 1.0))

iou loss map은 log를 사용하여 구하는데, 여기서 최종적으로 nan이 반환되는 것을 확인하였고 중점적으로 해결하고자 했습니다.

 

-> 이유가 뭐지? 🚨

-log 함수는 입력 값이 0이면 inf를 반환하고, 음수이면 nan을 반환합니다. 즉, 위에서는 (area_intersect + 1.0) / (area_union + 1.0)의 결괏값이 0이거나 음수이면 문제가 나타날 수 있습니다.

 

area_intersect와 area_union 이 동일한 값을 가지고 있을 때, 즉 bbox와 영역이 완벽히 겹치는 경우 area_에서 오버플로우가 발생하여 입력이 0이 되어 무한대로 발산할 수 있습니다.

 

-> AMP가 어떻게 동작하기에 이러한 문제가 발생한 것인지 ✅

 

https://hsb422.tistory.com/entry/MLOps-PART%ED%95%99%EC%8A%B5-%EC%86%8D%EB%8F%84-%EA%B0%9C%EC%84%A0

MLOps PART.학습 속도 개선

이전에 포스팅한 내용을 보면 AMP는 학습 속도를 높이고 메모리 사용량을 줄이기 위해 f16 연산을 사용합니다. 이는 f32보다 더 적은 범위의 값을 사용하여 오버, 언더플로우가 더 쉽게 발생할 수 있습니다.

 

-> 이슈 해결

iou_loss_map = -torch.log((area_intersect + 1.0) / (area_union + 1.0) + 6e-8)

따라서 매우 작은 값 6e-8(fp16의 최소값)을 더하여 log 결과가 발산하지 않도록 수정했습니다.

 

결과적으로 NaN 문제를 해결할 수 있었고 그 결과 40초의 성능 개선이 있었습니다.

 

3. 피클 적용

[고성능 ML 백엔드 피클링]

고성능 ML 백엔드를 위한 10가지 Python 성능 최적화 팁

위 하이퍼 커넥트 사례에 따르면 데이터 직렬화로 데이터 전달에서 성능을 높일 수 있다고 하여 적용해 보기로 했습니다.

 

-> 기존 데이터 셋

class SceneTextDataset(Dataset):
    def __init__(self):
        with open(osp.join(root_dir, 'ufo/{}.json'.format(split)), 'r') as f:
            anno = json.load(f)

    def __getitem__(self, idx):
            vertices, labels = [], []

init에서 json을 로드하고 getitem 부분에서 json을 바탕으로 이미지를 로드한 후 vertices, labels를 나누는 작업을 학습 과정에서 진행합니다. 학습을 진행하는 과정에서 배치 사이즈만큼 데이터 처리를 하다보니 시간이 매우 오래걸립니다.

 

-> pickle로 미리 해두기

vertices, labels = np.array(vertices, dtype=np.float32), np.array(labels, dtype=np.int64)
vertices, labels = filter_vertices(
    vertices,
    labels,
    ignore_under=ignore_under_threshold,
    drop_under=drop_under_threshold,
)

image = Image.open(image_fpath)

image, vertices = resize_img(image, vertices, image_size)
image, vertices = adjust_height(image, vertices)
image, vertices = rotate_img(image, vertices)
image, vertices = crop_img(image, vertices, labels, crop_size)

if image.mode != 'RGB':
    image = image.convert('RGB')

total["images"].append(np.array(image))
total["vertices"].append(vertices)
total["labels"].append(labels)

with open(osp.splitext(json_dir)[0] + ".pkl", "wb") as fw:
    pickle.dump(total, fw)

학습 단계의 데이터를 가져오는 과정에서 image, labels를 나누는 작업을 하지 않고 미리 다 처리를 한 후 pickle로 저장해서 학습 단계에서는 처리된 pickle 데이터만 가져오도록 구현하였습니다.

 

학습 데이터를 pickle로 저장하고,

 

데이터 로딩에서 pickle을 로드하도록 변경하였습니다.

 

그 결과 10초의 성능 개선이 있었습니다.

 

4. 버전 이슈

대회 커뮤니티에서 넘파이 버전에 따라 학습 속도가 대폭 감소한다는 정보를 공유받았습니다. 어떻게 낮은 버전을 사용했다는 것 만으로 학습 속도가 크게 차이나는 걸까요?

 

GPT :
1) 최적화 개선: 라이브러리의 업데이트는 종종 성능 최적화를 포함하며, 이는 학습 속도에 직접적인 영향을 미칠 수 있습니다. numpy나 opencv-python 등의 라이브러리는 이미지 처리 및 수치 계산에 사용되는 라이브러리로서, 이들의 성능 최적화는 학습 속도에 큰 영향을 미칠 수 있습니다.

2) 호환성 문제 해결: 특정 버전의 라이브러리들이 서로 잘 호환되지 않으면, 예기치 않은 성능 저하를 일으킬 수 있습니다. 버전을 낮춤으로써 이러한 호환성 문제가 해결되었을 수 있습니다.

ML은 수치적인 계산을 하는 라이브러리를 사용하기에 최적의 버전을 사용하는 것이 중요할 수도 있다는 고민을 하게 되었습니다. ML에서 파이썬을 대표 언어로 사용하는 것도 빠른 행렬 연산을 제공하기 때문이므로, 반드시 인지하고 있어야겠습니다.

 

결론

원본 > AMP > 피클
AMP : 약 40초 개선
피클 : 약 10초 개선
총 약 1분 개선

1 에폭당 약 1분이라는 속도 개선을 이뤘습니다. 전체 150 에폭이라면 2시간 30분을 절약할 수 있습니다. ML 개발 환경에서 직면할 수 있는 문제를 해결하는 경험을 할 수 있어서 뿌듯합니다.