Pytorch Resource & 모델 구조 Profiler 도구 (torch profiler)

Introduction

• 딥러닝 학습을 잘(?)한다는 것을 정의하기는 어렵지만, 더 빠른 시간 안에 많은 양을 학습하는 것은 매우 중요하다.
• 딥러닝의 모델은 다수의 layer로 구성되어 있기 때문에, 각 layer의 결과가 데이터가 어느 형태로 존재하는지, 어느 layer가 병목 현상인지를 파악하는 것이 까다롭다.
• 특히, 최근에는 모델을 직접 코드로 구현하기보다는,  pre-trained model을 사용하는 경우가 많은데, 사전에 사용하는 모델의 구조를 알지 못하면, 내부 동작을 제대로 파악할 수 없다.
• Pytorch에서는 이러한 모델의 구조와 각 layer에서의 cost를 profiling 할 수 있는 torch profiler를 지원한다.

 

Code Sample

• torch profiler 테스트를 위한 resnet18을 이용한 CIFAR-10 classification code이다. 
• 모델을 직접 코드로 구현한 것이 아닌, torchvision에서 load 하였다.
• 사용 상황을 가정하자면, load 한 모델의 구조를 모르거나, 모델에 부하가 존재하는 부분을 tuning 해야 하는데, 어느 layer를 바꿔야 할지 모르는 상황이다. 

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

TRAIN_EPOCH = 10
TRAIN_PRINT_FREQUENCY = 200

if __name__ == '__main__':
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomCrop(32, padding=4),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

    model = torchvision.models.resnet18(pretrained=False)
    num_features = model.fc.in_features
    model.fc = nn.Linear(num_features, 10)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    for epoch in range(TRAIN_EPOCH):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data[0].to(device), data[1].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % TRAIN_PRINT_FREQUENCY == TRAIN_PRINT_FREQUENCY - 1:
                print(f'Epoch: {epoch + 1}, Batch: {i + 1}, Loss: {running_loss / 200:.3f}')
                running_loss = 0.0
    print("Training finished.")

    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    print(f'Test Accuracy: {accuracy:.2%}')

 

Torch Profiler

[Setup]

• Pytorch에서는 1.8 버전 이상부터 torch의 profiling을 위한 torch.profiler를 제공한다. 따라서, Torch 버전이 1.8 이상인 경우에는 별도의 설치가 필요 없다. 
• 현재 설치된 Torch 버전을 잘 모른다면, Python에서 아래 명령어를 통해 확인해 보자.

import torch
print(torch.__version__)

• 만약, torch 버전이 1.8 미만에 torch 버전을 바꿔도 문제가 없는 상황이라면, 아래 명령어를 통해 torch 버전을 업그레이드해 준다. 

pip install --upgrade torch torchvision

 

[사용법]

• 사용방법은 매우 간단하다. 우선 torch.profiler를 import 하고, 콘텍스트 관리자(with 절)를 이용하여, profiling을 위한 부분을 감싸주면 된다. (함수 전체에 대한 profiling은 profile with 절을 @profile(acitivities~)와 같은 decorator로 처리할 수 있다.)
• 아래는 sample code 중, train에 대한 profiling을 위한 소스이다. 주의할 점은, profiling에 memory가 많이 소모되기 때문에, train epoch을 1로 낮춰놓고 profiling을 하는 것이 좋다. (어차피, 같은 동작이 반복되기 때문에, input 하나만을 측정해도 별 문제는 없다.)
• 모델을 GPU에서 돌려서  "ProfilerActivity.CUDA "를 포함했지만, CPU로 돌리는 환경에서는 해당 인자를 생략해도 된다. (다만, GPU 환경에서는 CPU, CUDA 인자 모두 필요함)

...
from torch.profiler import profile, record_function, ProfilerActivity
...

    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
        for epoch in range(TRAIN_EPOCH):
            running_loss = 0.0
            for i, data in enumerate(trainloader, 0):
                inputs, labels = data[0].to(device), data[1].to(device)
                optimizer.zero_grad()
                with record_function("model_train"):
                    outputs = model(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

                running_loss += loss.item()
                if i % TRAIN_PRINT_FREQUENCY == TRAIN_PRINT_FREQUENCY - 1:
                    print(f'Epoch: {epoch + 1}, Batch: {i + 1}, Loss: {running_loss / 200:.3f}')
                    running_loss = 0.0
    print(prof.key_averages().table(sort_by="self_cpu_time_total"))
    print("Training finished.")

[인자 설명]

• activities : list 형태로 입력받는다. 어떤 활동을 profiling 할 것인지를 지정한다. 가능한 활동은 다음과 같다.
  • ProfilerActivity.CPU : CPU 작업(연산, 함수 호출)에 대한 프로파일링, CPU 시간, 메모리 사용량등을 제공
  • ProfilerActivity.CUDA  : CUDA 작업(GPU 연산, 호출)에 대한 프로파일링, GPU 시간, 메모리 사용량등을 제공
• record_shapes : bool 형태, 각 layer의 입력(input)을 기록할지 여부
• profile_memory : bool 형태, memory를 profiling 할지 여부, False로 설정하면 time에 대한 profiling만 진행한다.
• on_trace_ready : Profiling 결과가 준비되었을 때, 호출될 callback 함수를 지정할 수 있음. on_trace_ready 옵션을 통해, 함수를 사전 정의해, profiling 결과 등을 file 형태로 떨굴 수 있다.

  ...
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True,
                 profile_memory=True, on_trace_ready=finish_profiler) as prof:
  ...

• with_stack : bool 형태, 함수 호출 stack 정보를 표기할지에 대한 여
• with_flops : bool 형태, 실제로 계산 비용을 FLOPs로 측정한 결과 
• with_modules : bool 형태, profiling 결과에 연산의 호출 stack에 대한 module의 계층 구조를 기록해 줌. (어떤 연산이 어떤 연산의 내부에서 호출되었는지를 나타내줌)

 

 

[결과 출력]

• 결과는 다음과 같은 명령어로 호출할 수 있다.

print(prof.key_averages().table(sort_by="self_cpu_time_total"))

• 결과는 table 형태로 보이는데, 아래와 같이 다양한 옵션들을 사용할 수 있다. (key_averages() 내의 인자 형태로 들어감)
  • group_by_input_shape : True로 설정하면, 동일한 입력 모양을 가진 연산 또는, 함수 호출을 grouping 할 수 있다. (모델의 input 사이즈를 보려면, 해당 옵션을 True 설정해야 한다.)
  • group_by_stack_n : 연산 또는 함수의 stack의 상위 n 단계만을 기준으로 grouping 할지 지정하는 인자
• table의 인자도 지정할 수 있는데, table의 출력에 대한 옵션을 지정한다.
  • sort_by : table을 어떤 기준으로 order by 할지 (default : None)
  • row_limit : 몇 개까지 표시할지
  • header : header를 표시할지 (default : None)
  • top_level_events_only : 해당 옵션을 True 설정하면, 최상위 호출 단계까지만 표시
• 아래는 sample code에 대해 profiling을 수행한 결과이다. (CPU 시간이 큰 10개만 추출)

• 결과에서 보이는 각 칼럼은 다음과 같다.  
• CPU Time 관련 
  • Self CPU % : 연산 or 함수 호출이 소비한 CPU 시간의 백분율 (전체 실행 시간에서 해당 연산이 소요한 CPU 시간)
  • Self CPU : 해당 연산 or 함수 호출이 소비한 총 CPU 시간
  • CPU  total % : 해당 연산 or  함수 호출과 그 하위 호출에서 소요된 총 CPU 시간의 백분율
  • CPU total : 해당 연산 or 함수 호출과 그 하위 호출에 의해 사용된 총 CPU 시간
  • CPU time avg : 해당 연산 or 함수 호출의 평균 CPU 시간 (평균적으로 해당 연산이 소요되는 시간)
• CUDA Time 관련  
  • Self CUDA : 해당 연산 or 함수 호출이 소비한 총 CUDA 시간
  • Self CUDA % : 해당 연산 or 함수 호출이 소비한 총 CUDA 시간의 백분율
  • CUDA total : 해당 연산 or 함수호출과 그 하위 호출에서 소요된 총 CUDA 시간
  • CUDA time avg : 해당 연산 or 함수호출과 그 하위 호출에서 소요된 평균 CUDA 시간
  • # of Calls : 해당 연산 또는 함수 호출의 호출 횟수
• Model Input 관련 
  • Input Shapes : record shapes를 True로 하고, key_averages에 group_by_input_shape를 true로 지정한 경우에만 보인다. 각 연산의 input shape이 보인다.
• CPU memory 관련 (snapshot 형태기 때문에 사용 전과 후의 memory 사용 delta값이 나온다. 즉, 음수가 될 수 있다.) 
  • CPU Mem : 연산 or 함수 호출이 소비한 CPU의 메모리 총 용량 
  • Self CPU Mem : 연산 or 함수 호출이 직접적으로 사용한 CPU 메모리 용량
• CUDA memory 관련 (snapshot 형태기 때문에 사용전과 후의 memory 사용 delta값이 나온다. 즉, 음수가 될 수 있다.)
  • CUDAMem : 연산 or 함수 호출이 소비한 CUDA의 메모리 총 용량 
  • Self CUDAMem : 연산 or 함수 호출이 직접적으로 사용한 CPU 메모리 용량
• 연산량 관련
  • Total MFLOPs : 연산 or 함수 호출이 실행될 때, 총 수행된 MFLOPs 수 

 

Torch 모델에서 torch Profiling을 통해, 부하가 되는 부분이나, Layer의 input size 등을 확인할 수 있다. 해당 profiling은 모델에서 부하가 되는 부분을 개선하거나, 하드웨어 확장에 대한 의사결정, Batch size 조절 등 다양한 model 개선에 사용될 수 있다.