[최적화] pytorch CSR 구현 (작성중)

Pruning 이후, 파라미터는 0이 많은 Sparse matrix(희소 행렬)의 형태를 갖게 된다. 하지만, 이 상태는 파라미터 값이 0으로 설정되었을 뿐, 사실상 경량화에 직접적으로 도움이 되지 않는다. 따라서 모델의 용량을 줄이기 위해서는 이 sparse matrix 를 효율적으로 표현할 필요가 있다. 이를 위해 제안된 방법으로 COO(Coordinate Format), CSR(Compressed Sparse Row), CSR(Compressed Sparse Column) 등이 있다. 본 글에서는 이를 설명하고, pytorch를 활용하여 구현해보고자 한다.

 

 

COO이란

: 0이 아닌 값의 값, 행, 열을 별도의 메모리에 저장하는 방법 (3a개 메모리 필요,a=0이 아닌 개수)

위의 예시에서는 기존 12개 메모리가 필요했으나 COO를 활용한다면 9개의 메모리만 사용한다.

 

CSR이란

0이 아닌 값의 값과 열을 별도의 메모리에 저장하고, 행은 포인터(주소값)형태로 가지고 있는 방법 (2a+(n+1))개 메모리 필요, n=행의 길이)

*CSC는 CSR과 동일 한 개념으로, 행과 열만 바뀐 방법.

 

본 글 외에도 희소 행렬 표현에 대한 설명 글 혹은 구현해 놓은 글을 쉽게 찾아볼 수 있었다. 하지만, 모두 희소 행렬 표현 구현까지만 하였을 뿐, 실제로 어떠한 방식으로 모델에서 활용되는지는 확인할 수 없었다. 본 글에서는 희소 행렬을 사용하여, 모델 저장, 추론 등 pytorch에서 실질적으로 이득을 보는 것을 확인해보고자 한다.

 

간단하게 행렬곱에서 희소 행렬을 사용했을 때와 사용하지 않았을 때의 시간을 비교해보았다.

1. 희소 행렬을 사용하지 않은 경우

# 희소 Tensor 생성
dense_tensor = torch.Tensor([[0, 0, 3], [4, 0, 5], [0, 0, 0]])

s=time.time()
# 밀집 Tensor와의 행렬 곱셈
dense_matrix = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = torch.mm(dense_tensor, dense_matrix)
e=time.time()

print(f"{e - s:.5f} sec")

print("Dense Tensor:")
print(dense_tensor)

print("Result of Matrix Multiplication:")
print(result)

 

2. 희소 행렬을 사용한 경우

import torch
import time

# 희소 텐서 생성
indices = torch.LongTensor([[0, 1, 1], [2, 0, 2]])
values = torch.FloatTensor([3, 4, 5])
shape = torch.Size([3, 3])
sparse_tensor = torch.sparse.FloatTensor(indices, values, shape)

# 밀집 Tensor로 변환
dense_tensor = sparse_tensor.to_dense()

s=time.time()
# 밀집 Tensor와의 행렬 곱셈
dense_matrix = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = torch.mm(dense_tensor, dense_matrix)
e=time.time()

print(f"{e - s:.5f} sec")

# COO format
print("Sparse Tensor (COO):")
print(sparse_tensor)

print("Dense Tensor:")
print(dense_tensor)

print("Result of Matrix Multiplication:")
print(result)

오히려 COO format을 사용한 경우에 시간이 더 많이 걸리는 것을 확인할 수 있었다. 아마 torch.mm이 희소 행렬 computing 지원 연산이 아니라, 텐서로 변환 후에 수행하다 보니, 시간이 더 오래 걸리는 것 같다. 

이 내용은 다시 수정해야 될 것 같다 (pytorch에서 sparse compution 지원 연산이 잇는지 찾아봐야 함)

import torch
import time
from scipy.sparse import random

# 희소 행렬 생성 (SciPy 이용)
sparse_matrix = random(5000, 5000, density=0.01, format='coo')

# 밀집 행렬 생성
dense_matrix = torch.randn(5000, 5000)

# 희소 행렬에서의 곱셈
start_time = time.time()
sparse_result = torch.mm(torch.sparse.FloatTensor(torch.LongTensor([sparse_matrix.row, sparse_matrix.col]), torch.FloatTensor(sparse_matrix.data), torch.Size(sparse_matrix.shape)), dense_matrix)
sparse_elapsed_time = time.time() - start_time

# 밀집 행렬에서의 곱셈
start_time = time.time()
dense_result = torch.mm(dense_matrix, dense_matrix)
dense_elapsed_time = time.time() - start_time

print("Sparse Computing - Elapsed Time: {:.4f} seconds".format(sparse_elapsed_time))
print("Dense Computing - Elapsed Time: {:.4f} seconds".format(dense_elapsed_time))

# 결과 비교 (에러 확인)
error = torch.norm(sparse_result - dense_result)
print("Error between Sparse and Dense Results:", error)

희소행렬 X 밀집행렬이 밀집행렬 X 밀집행렬이 더 빠르게 수행된다. 저장 방식과는 관계없이 0값을 무시하고 연산이 수행되기 때문에 빠르게 수행된다. (pytorch 코드 확인해볼 필요 있음)

 

pytorch 모델로는 가장 기본 모델인 LeNet을 사용한다. 

 

 

 

--> 이거를 실제 pytorch에서 이득을 보는 방법?

-> 학습 할때, 결국, 희소 행렬로 만들어줘야 하는거 아닌가???

-> 모델 저장, 학습, 추론에서 직접적으로 이득을 볼 수 있는지 확인이 필요함 그래야 pruning이 경량화 기법이라고 납득할 수 있을 것 같음