옵티마이저 크기 계산하기

모델 학습에 사용할 여러 옵티마이저를 비교하기 위해 메모리 사용량을 정량적으로 측정하려고 합니다. 테스트를 위해 DistilBERT 모델과 AdamW 옵티마이저를 불러왔으며, 이제 메모리 사용량을 계산해 보겠습니다. 옵티마이저의 크기를 계산하는 compute_optimizer_size 함수를 작성해 보세요.

AdamW optimizer는 Trainer 없이 직접 정의되었으며, 학습이 완료된 상태입니다.

이 연습은 강의의 일부입니다

PyTorch로 AI 모델 효율적으로 학습시키기

강의 보기

연습 안내

for 루프 안에서 각 tensor의 요소 수와 요소 크기를 계산하세요.
optimizer의 전체 크기를 메가바이트 단위로 계산하세요.
optimizer.state의 적절한 메서드를 사용하여 옵티마이저 상태 딕셔너리에 접근하세요.

실습형 인터랙티브 연습

이 예제를 이 샘플 코드를 완성하여 풀어보세요.

def compute_optimizer_size(optimizer_state):
    total_size_megabytes, total_num_elements = 0, 0
    for params in optimizer_state:
        for name, tensor in params.items():
            tensor = torch.tensor(tensor)
            # Compute number of elements and size of each tensor
            num_elements, element_size = tensor.____(), tensor.____()
            total_num_elements += num_elements
            # Compute the total size in megabytes
            total_size_megabytes += ____ * ____ / (1024 ** 2)
    return total_size_megabytes, total_num_elements

# Pass in the optimizer state
total_size_megabytes, total_num_elements = compute_optimizer_size(optimizer.state.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")

코드 편집 및 실행