# [기술] 학습시 메모리사용량 감소

안규원·2024년 6월 10일

AI

목록 보기
9/22

[Quantization]


양자화는 모델에서 숫자를 나타내는 비트 수를 줄이는 과정이다.

import torch
import torchvision

# 사전 훈련된 모델 불러오기
model = torchvision.models.resnet18(pretrained=True)

# 모델 융합
model_fused = torch.quantization.fuse_modules(model, [['conv1', 'bn1', 'relu']])

# 양자화 구성 지정
model.qconfig = torch.quantization.default_qconfig

# 정적 양자화를 위한 모델 준비
torch.quantization.prepare(model, inplace=True)

# 대표 데이터를 사용하여 모델 보정
# `calibration_data`가 보정 데이터셋을 가진 DataLoader라고 가정
with torch.no_grad():
    for data, _ in calibration_data:
        model(data)

# 양자화된 버전으로 모델 변환
torch.quantization.convert(model, inplace=True)

# 양자화된 모델 저장
torch.save(model.state_dict(), 'quantized_model.pth')

[Flash Attention]


Transformer 모델에서 Attention의 효율성을 높이는 방법으로, 연산 순서를 재정렬하고 중복 계산을 줄여 메모리 사용량을 줄인다.

[PEFT]


모델 파라미터의 일부만 미세 조정하거나 효율적인 미세 조정을 위해 추가 파라미터를 도입하는 방법으로, 제한된 컴퓨팅 자원으로 특정 작업에 맞게 조정한다.

import torch
import torch.nn as nn
from transformers import BertModel

class Adapter(nn.Module):
    def __init__(self, input_dim, adapter_dim):
        super(Adapter, self).__init__()
        self.linear_down = nn.Linear(input_dim, adapter_dim)
        self.relu = nn.ReLU()
        self.linear_up = nn.Linear(adapter_dim, input_dim)
    
    def forward(self, x):
        down = self.linear_down(x)
        activation = self.relu(down)
        up = self.linear_up(activation)
        return up + x  # 잔차 연결

# 사전 훈련된 BERT 모델 불러오기
model = BertModel.from_pretrained('bert-base-uncased')

# 모델에 어댑터 삽입
adapter_dim = 64
for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        module.adapter = Adapter(module.out_features, adapter_dim)

# 어댑터와 함께 모델 미세 조정
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fn = nn.CrossEntropyLoss()

# 예제 훈련 루프
for batch in training_data:
    inputs, targets = batch
    outputs = model(**inputs)
    loss = loss_fn(outputs.logits, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

0개의 댓글