Pytorch Profiler Tensorboard로 시각화

Introduction

• Pytorch 학습 중, Resource와 모델 구조에 대한 profiling은 torch profiler를 이용해 가능하였다.

2023.07.09 - [Python] - Pytorch 구조 & Resource Profiler 도구 (torch profiler)

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

• 이때, profiling의 결과는 테이블 구조의 텍스트 형태로 터미널에 출력 or 파일에 저장 가능하다.
• 이 결과로도 Insight를 충분히 추출할 수 있지만, 텍스트 형태로 분석하다 보면, 가시성이 떨어진다는 점과, profiling 이력 간 비교가 어렵다는 점이 아쉽다.
• 이를 해결할 수 있는 딥러닝의 시각화 툴인 Tensorboard에 profiling 결과를 올리는 방법이 있어서, 이 방법을 알아보고자 한다.

 

원리

• torch profiler의 옵션에 "on_trace_ready"라는 옵션이 존재한다. 이것은 profiling 결과가 준비되었을 때, 호출할 callback 함수를 지정하는 것인데, 이 callback 함수로 tensorboard에서 제공하는 trace handler를 연결하여, tensorboard에서 읽을 수 있는 log 형태로 떨궈준다.

 

Setup

• tensorboard에서 torch profiler의 결과를 읽어서 표현할 수 있는 plugin을 추가로 설치해야한다. (당연히, tensorboard가 필요하기 때문에, 아래 plugin을 설치하면, tensorboard도 자동으로 설치된다.)

pip install torch_tb_profiler

 

사용법

• torch profiler를 이용하기 위한, 콘텍스트 관리자(with 절)에  on_trace_ready 옵션에, "tensorboard_trace_handler" 함수를 지정해 준다.
• 이때, tensorboard_trace_handler 함수의 인자로, tensorboard에서 읽을 수 있도록 log 디렉터리를 지정해 준다.

  ...
  with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True,
                 profile_memory=True, on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18')) 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_inference"):
                    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
  ...

• 그 후에, tensorboard 실행을 해주면 끝난다.

tensorboard --logdir=./log/resnet18

 

결과

[NORMAL]

• Overview : Overview는 전체적인 프로파일링의 결과에 대해 보여주는 화면이다.
  • Configuration : Profiling 시 사용된 설정 정보를 표시한다. 
  • Execution Summary : 전체 수행 시간과, 각 단계에 소요된 수행시간을 보여준다.
  • Spen Time Breakdown : 코드 또는 연산의 실행 시간을 단계별(Kernel, Memcpy, Memset, Runtime, DataLoader, CPU Exec, Other)로 분해하여 보여준다. 
  • Performance Recommentation : profiling 결과를 기반으로 한 성능 개선 권장 사항을 자동으로 생성해 준다. (실제로 어떤 원리로 동작하는지는 잘 모른다.)

• Operator : 연산에 대한 profiling 결과를 보여준다. torch profiling 결과를 터미널에서 출력하였을 때, 보여주던 결과를 그대로 보여준다고 생각하면 된다. 추가적으로 Tensor Cores Eligible 옵션이 있는데, 해당 연산이 GPU를 사용할 수 있는지에 대한 가능 여부를 표시한 것이다. Group By 조건을 바꾸면, input shape도 볼 수 있다.

 

• Trace : 함수 및 연산의 실행 시간을 시간 경과에 따라 그래프 형태로 표시한다. 해당 코드의 실행에 사용된 Process와 그 안의 Thread의 동작을 확인 할 수 있다.  (사실 이 UI는 torch profiler의 chrome tracing 기능으로도 볼 수 있다.)

 

• Memory : 실행 시간에 따른 Memory 사용 추이를 보여준다. 각 연산마다 할당한 메모리와, Allocation Time과 Release Time, Duration을 보여준다. 코드 실행에 사용한 H/W 별로 볼 수 있다. (다만, 해당 UI에서 Memory를 많이 사용하는지, Chrome이 계속 죽는다.)

 

 

[DIFF]

• Tensorboard를 통한 torch profiler 시각화의 가장 큰 장점이라고 할 수 있는 이력 간 비교 기능이다. Baseline의 log를 정한 뒤, 비교하고자 하는 log를 대입하면, 그 둘 간의 profiling 결과의 delta 값을 보여준다. 
• 이를 통해, Profiling 결과를 비교하면서, H/W 효율화를 위한 구조 개선을 진행할 수 있다.

 

 

Torch 모델에서 torch profiling을 시각화하여 비교하는 방법을 알아보았다. 개인 프로젝트에서는 그 효용이 덜하겠지만, 많은 사람들이 같은 모델을 연구할 때, 성능과 profiling 결과를 모두 tensorboard를 통해 시각화하여, 성능을 유지하면서 모델의 연산 효율성을 향상하거나, 모델의 연산  효율성을 유지하면서 모델의 성능을 향상하는 데, 사용하면 매우 유용할 것이다.