[INSIGHT] 딥러닝에서 Tensor로 연산 효율 개선하기

2025.01.04

 

딥러닝을 위한 모델의 코드에서 시간 및 메모리 효율을 높이기 위해 주로 Tensor 연산을 사용한다.

 

인공지능 모델의 코드를 Tensor 연산을 제대로 고려하지 않은 채, 평소에 프로그래밍 언어를 배우던 방식으로 코드를 작성하였다.

그러다 보니, Tensor 연산의 이점을 간과하고 상당히 비효율적인 for문을 사용하였다.

 

이를 개선하였더니, 기존의 모델과 비교하여 70% 가량의 실행 시간을 단축할 수 있었다.

 

왜 딥러닝에서 Tensor 연산을 사용할까?

1. 연산 동시 수행 및 병렬 처리 : 텐서 연산은 GPU에서 벡터화 연산으로 처리되기에 여러 연산이 동시에 수행

2. 메모리 효율 : 불필요한 데이터 복사 및 추가적인 메모리 사용 최소화

 

Tensor 연산화 과정

• Cross Attention 모듈 
  • 유전자와 약물의 하위 구조 간의 가중치를 패스웨이 마다 학습하는 과정
  • 텐서 연산으로 한번에 병렬적으로 처리 가능한 부분이었으나, for문을 통한 비효율적인 연산을 진행

# 적용 전) Cross Attention 모듈 - v2_main_dummy.ipynb
updated_gene_embeddings = []
updated_drug_embeddings = []

for i in range(gene_embeddings.size(1)):
    gene_emb = gene_embeddings[:, i, :]  
    sub_emb = substructure_embeddings[:, i, :]  

    updated_gene_emb = self.Gene2Sub_cross_attention(gene_emb, sub_emb)
    updated_sub_emb = self.Sub2Gene_cross_attention(sub_emb, gene_emb)

    updated_gene_embeddings.append(updated_gene_emb)
    updated_drug_embeddings.append(updated_sub_emb)

updated_gene_embeddings = torch.stack(updated_gene_embeddings, dim=1)


# 적용 후) Cross Attention 모듈 - v3_main_dummy.ipynb
gene_query = self.Gene2Sub_cross_attention(gene_embeddings, substructure_embeddings)
sub_query = self.Sub2Gene_cross_attention(substructure_embeddings, gene_embeddings)

 

• Pathway Graph Embedding 모듈
  • Cross Attention 모듈 후 임베딩을 패스웨이 마다 그래프의 노드피처를 업데이트 후 그래프 신경망을 통해 학습하는 과정
  • 메모리 비효율적인 append 함수를 대체하여 리스트 생성 후 즉시 stack 함수를 활용하여 불필요한 메모리 할당 제거

# 적용 전) Pathway Graph Embedding 모듈 - v2_main_dummy.ipynb
pathway_graph_embeddings = []
for pathway_idx in range(updated_gene_embeddings.size(1)):
    gene_emb = updated_gene_embeddings[:, pathway_idx, :, :]
    pathway_graph = pathway_graphs[pathway_idx]
    pathway_graph_embedding = self.pathway_graph(gene_emb, pathway_graph)
    pathway_graph_embeddings.append(pathway_graph_embedding)

pathway_graph_embeddings = torch.stack(pathway_graph_embeddings, dim=1)


# 적용 후) Pathway Graph Embedding 모듈 - v3_main_dummy.ipynb
pathway_graph_embeddings = torch.stack([
    self.pathway_graph(gene_query[:, i, :, :], pathway_graphs[i])
    for i in range(gene_query.size(1))
], dim=1)

 

 

Tensor 연산화 후 결과 비교

모델	실행 시간
최적화 전 (Tensor 연산화 이전)	1 min 40.0 sec
최적화 후 (Tensor 연산화 이후)	31.9 sec

 

최적화 전, 후 실행 결과 비교