작은 행렬 연산 벤치마크 테스트(파이토치, 러스트)

rust candle github의 이슈 댓글

 

최근 Rust언어로 작성된 Candle이라는 라이브러리에 관심을 가지며 파이토치와 Candle로 작성한 코드의 벤치를 간단하게 테스트 해보던 중에, 다음과 같은 요청을 받았다.

본인의 도메인에서는 작은 행렬을 연산하는 것이 일반적인데, 파이토치가 여기에 bad job이라고 한다.

실제로 확인하기 위해서 코드를 작성 후 테스트를 하였다

 

테스트 환경

CPU : AMD Eyzen Threadripper PRO 3955WX 16-cores

GPU : A6000

python version: 3.11.5

pytorch version: 2.1.1

numpy version: 1.24

rust version: 1.74.0

candle version: 0.3.1

 

테스트

 

파이토치, CPU

3 x 3 행렬을 각각 1천만개씩 만들어 CPU로 연산

import torch
import time
from tqdm import tqdm

# CPU에서 텐서 생성
tensor_cpu = torch.randn(10000000,3,3)
tensor_cpu2 = torch.randn(10000000,3,3)

# 연산 시작 시간 기록
start_time = time.time()

# 1천만번 연산 수행
for a, b in tqdm(zip(tensor_cpu, tensor_cpu2), total=tensor_cpu.shape[0]):
    result_cpu = torch.mm(a,b)

# 연산 종료 시간 기록 및 소요 시간 계산
end_time = time.time()
elapsed_time = end_time - start_time

print("CPU에서 1천만번 연산 수행 시간: {}초".format(elapsed_time))

## 55.352723초

 

 

파이토치, CPU

위와 비슷하나, torch.mm 메서드가 아닌 matmul로 테스트

numpy의 경우 dot 연산과 matmul 연산이 달라 확인 차원에서 다음과 같이 테스트 하였다.

import torch
import time
from tqdm import tqdm

# CPU에서 텐서 생성
tensor_cpu = torch.randn(10000000,3,3)
tensor_cpu2 = torch.randn(10000000,3,3)

# 연산 시작 시간 기록
start_time = time.time()

# 1천만번 연산 수행
for a, b in tqdm(zip(tensor_cpu, tensor_cpu2), total=tensor_cpu.shape[0]):
    result_cpu = torch.matmul(a,b)

# 연산 종료 시간 기록 및 소요 시간 계산
end_time = time.time()
elapsed_time = end_time - start_time

print("CPU에서 1천만번 연산 수행 시간: {}초".format(elapsed_time))
55.52095484초

 

파이토치, GPU

파이토치, torch.mm 메서드를 활용, GPU 연산

위 사람이 요청할 때 언급했던 대로, GPU 연산에서 보다 높은 시간이 소요됨을 확인하였다.

# GPU 사용 가능 여부 확인
if torch.cuda.is_available():
    # GPU에서 텐서 생성
    tensor_gpu = torch.randn(10000000,3, 3, device='cuda')
    tensor_gpu2 = torch.randn(10000000,3, 3, device='cuda')

    # 연산 시작 시간 기록
    start_time = time.time()

    # 1천만번 연산 수행
    for a, b in tqdm(zip(tensor_gpu, tensor_gpu2), total=tensor_gpu.shape[0]):
        result_gpu = torch.mm(a,b)

    # 연산 종료 시간 기록 및 소요 시간 계산
    end_time = time.time()
    elapsed_time = end_time - start_time

    print("GPU에서 1천만번 연산 수행 시간: {}초".format(elapsed_time))

## 112.187461초

 

 

넘파이, 일반 cpu 연산

넘파이 또한 수치 연산에서 가장 일반적으로 사용되는 라이브러리로, 토치와 비교를 하기 위해 테스트 하였다

테스트 결과, 14초로 토치에 비해 훨씬 빠르게 연산을 끝낸다

import numpy as np
import time
from tqdm import tqdm

numpy_tensor1 = np.random.rand(10000000, 3, 3)
numpy_tensor2 = np.random.rand(10000000, 3, 3)

start_time = time.time()

for a, b in tqdm(zip(numpy_tensor1, numpy_tensor2), total=numpy_tensor2.shape[0]):
    numpy_result = np.dot(a,b)
    end_time = time.time()
    elapsed_time = end_time - start_time

print("CPU에서 1천만번 연산 수행 시간: {}초".format(elapsed_time))

## 14.346451초

 

러스트로 넘어가기에 앞서,  파이썬에서 최대한 성능을 이끌어내고자, 1가지 추가 테스트를 수행했다

 

넘파이, 넘바 jit compile CPU

파이썬의 수치 연산에서, 고성능을 이끌어내기 위해 일반적으로 사용하는 라이브러리인 넘바를 사용했다.

병렬 처리 또한 할 수 있으나, 별도로 하지 않고 진행, CPU 연산 중 가장 고성능을 보여주었다.

from numba import jit
import numpy as np
from tqdm import tqdm

@jit(nopython=True)
def main(numpy_tensor1, numpy_tensor2):

    for a, b in zip(numpy_tensor1, numpy_tensor2):
        numpy_result = np.dot(a,b)
    return numpy_result

numpy_tensor1 = np.random.rand(10000000, 3, 3)
numpy_tensor2 = np.random.rand(10000000, 3, 3)
start_time = time.time()
main(numpy_tensor1=numpy_tensor1, numpy_tensor2=numpy_tensor2)
end_time = time.time()
elapsed_time = end_time - start_time
print("CPU에서 1천만번 연산 수행 시간: {}초".format(elapsed_time))

## 3.616512초

 

 

러스트

C++에 버금가는 언어로서 많은 영역에서 러스트 언어를 채택해 코드를 재작성하고 있으며, 특히 성능과 뛰어난 안정성 부분에서 주목받고 있는 언어이다.

필자는 특히 자연어처리 분야에서 매우 유명한 허깅페이스 기업에서, 러스트로 라이브러리를 만드는 작업 중에 있어 관심을 가지고 배우게 되었다.

 

Candle, CPU

파이토치의 API와 비슷하면서 동시에, 컴파일 용량을 줄이기 위해 최소한의 기능을 담고 있는 라이브러리이다.

늘어난 코드량에 비하면 미미한 차이이긴 하나, CPU로 극한의 성능을 보여주는 부분이다.

fn cpu_test() {
    let a = Tensor::randn(0.1f32, 1f32, (10000000,3,3), &Cpu);
    let b = Tensor::randn(0.1f32, 1f32, (10000000, 3,3), &Cpu);
    

    let start = Instant::now();
    let res: Vec<_> = zip(a,b)
    .into_iter()
    .map(move |(x, y)| {x.matmul(&y).unwrap();})
    .collect();
    
/*
    let start = Instant::now();
    for (arr1, arr2) in tqdm(zip(a, b),) {
        let res = arr1.matmul(&arr2);
    }
 */
    let end = start.elapsed().as_secs_f32();
    eprintln!("end = {:?}", end);
    eprintln!("res = {:?}", res);
}

## 2.7602초

 

Candle, GPU

Candle에선 일정 컴퓨팅 이상의 GPU에 한해서 CUDA 사용이 가능하다. 파이토치의 경우 GPU 사용시 오버헤드 발생으로 시간이 증가한데 반해서, 러스트에서는 안정적으로 고성능의 퍼포먼스를 보여주었다

fn gpu_test() {

    let device = Device::new_cuda(0).unwrap();

    let a = Tensor::randn(0.1f32, 1f32, (10000000,3,3), &device);
    let b = Tensor::randn(0.1f32, 1f32, (10000000, 3,3), &device);

    let start = Instant::now();
    let res: Vec<_> = zip(a,b)
    .into_iter()
    .map(move |(x, y)| {x.matmul(&y).unwrap();})
    .collect();
    let end = start.elapsed().as_secs_f32();
    eprintln!("end = {:?}", end);
    eprintln!("res = {:?}", res);
}
## 0.0005s