코드 리뷰

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models

# LoRA Layer
class LoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4):
        super(LoRALayer, self).__init__()
        self.rank = rank
        self.lora_a = nn.Parameter(torch.randn(in_features, rank))
        self.lora_b = nn.Parameter(torch.randn(rank, out_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))
        nn.init.zeros_(self.lora_b)

    def forward(self, x):
        return x + torch.matmul(torch.matmul(x, self.lora_a), self.lora_b)

# Modified ResNet with LoRA
class ResNetLoRA(nn.Module):
    def __init__(self, base_model, num_classes=10, lora_rank=4):
        super(ResNetLoRA, self).__init__()
        self.base_model = base_model
        self.lora_layers = nn.ModuleList([LoRALayer(layer.in_features, layer.out_features, lora_rank)
                                          for layer in base_model.modules() if isinstance(layer, nn.Linear)])
        self.fc = nn.Linear(base_model.fc.in_features, num_classes)
        self.base_model.fc = nn.Identity()  # Remove original fully connected layer

    def forward(self, x):
        x = self.base_model(x)
        for lora_layer in self.lora_layers:
            x = lora_layer(x)
        x = self.fc(x)
        return x

# Data preparation
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# Model, Loss, Optimizer
base_model = models.resnet18(pretrained=True)
model = ResNetLoRA(base_model, num_classes=10, lora_rank=4)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader)}')

print('Training Finished.')

import torch
import torch.nn as nn
import math

class LoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4):
        super(LoRALayer, self).__init__()
        self.rank = rank
        self.lora_a = nn.Parameter(torch.randn(in_features, rank))
        self.lora_b = nn.Parameter(torch.randn(rank, out_features))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))
        nn.init.zeros_(self.lora_b)

    def forward(self, x):
        return x + torch.matmul(torch.matmul(x, self.lora_a), self.lora_b)

from torchvision import models

class ResNetLoRA(nn.Module):
    def __init__(self, base_model, num_classes=10, lora_rank=4):
        super(ResNetLoRA, self).__init__()
        self.base_model = base_model
        self.lora_layers = nn.ModuleList([LoRALayer(layer.in_features, layer.out_features, lora_rank)
                                          for layer in base_model.modules() if isinstance(layer, nn.Linear)])
        self.fc = nn.Linear(base_model.fc.in_features, num_classes)
        self.base_model.fc = nn.Identity()  # 원래 fully connected 레이어 제거

    def forward(self, x):
        x = self.base_model(x)
        for lora_layer in self.lora_layers:
            x = lora_layer(x)
        x = self.fc(x)
        return x

# 기본 모델 로드
base_model = models.resnet18(pretrained=True)
model = ResNetLoRA(base_model, num_classes=10, lora_rank=4)

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader)}')

print('Training Finished.')

설명 요약

1. LoRALayer: 입력과 출력을 Low-Rank 행렬로 근사하여 추가 파라미터를 최소화합니다.

2. ResNetLoRA: ResNet-18 모델을 기반으로 LoRA 레이어를 추가하고, 원래의 fully connected 레이어를 대체합니다.

3. 데이터 준비: CIFAR-10 데이터셋을 로드하고 전처리합니다.

4. 손실 함수 및 최적화기 설정: 미리 훈련된 ResNet-18 모델을 LoRA와 결합하여 사용합니다.

5. 훈련 루프: 모델을 훈련시키고 각 epoch마다 손실을 출력합니다.