파이토치 (3) - 파이토치 구조 분석 & 구조별 문법

이영락·2024년 9월 13일

개발자 기본기

목록 보기
9/53

1. 데이터 준비 (Data Preparation)

딥러닝 모델 학습에 필요한 데이터를 준비합니다. PyTorch에서는 torch.utils.data.DatasetDataLoader를 사용하여 데이터를 불러오고 배치 단위로 처리합니다.

Dataset의 세 가지 필수 메서드

  • __init__: 이 메서드는 데이터셋 객체가 생성될 때 한 번만 실행되며, 데이터셋을 로드하는 데 필요한 초기 설정을 합니다.
  • __len__: 데이터셋의 샘플 개수를 반환하며, 이는 모델 학습에 필요한 데이터셋의 크기를 파악하는 데 사용됩니다.
  • __getitem__: 주어진 인덱스에 해당하는 샘플을 데이터셋에서 불러와 반환하는 메서드입니다. 데이터와 레이블을 반환하는 것이 일반적입니다.
import torch
from torch.utils.data import DataLoader, Dataset

# 사용자 정의 데이터셋 클래스
class loadDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 데이터셋 초기화
dataset = loadDataset(data, labels)
# DataLoader로 배치 단위로 데이터를 불러옴
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

2. 모델 정의 (Model Definition)

신경망 모델을 정의.
PyTorch에서는 torch.nn.Module을 상속받아 모델을 정의하고, forward() 함수에서 순전파를 정의합니다.

import torch.nn as nn
import torch.nn.functional as F

# 신경망 모델 정의
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(784, 128)  # 입력 크기: 784, 출력 크기: 128
        self.fc2 = nn.Linear(128, 10)   # 출력 크기: 10 (10개의 클래스 분류)

    def forward(self, x):
        x = F.relu(self.fc1(x))  # ReLU 활성화 함수 적용
        x = self.fc2(x)
        return x

# 모델 초기화
model = SimpleNN()

3. 손실 함수 및 옵티마이저 설정 (Loss Function & Optimizer)

모델 학습에 필요한 손실 함수와 옵티마이저를 설정합니다. 손실 함수는 모델의 성능을 평가하는 데 사용되고, 옵티마이저는 모델의 매개변수를 업데이트합니다.

import torch.optim as optim

# 손실 함수 설정 (교차 엔트로피 손실)
criterion = nn.CrossEntropyLoss()

# 옵티마이저 설정 (Adam 옵티마이저)
optimizer = optim.Adam(model.parameters(), lr=0.001)

4. 장치 설정 (Device Setting)

PyTorch에서는 GPU 또는 CPU를 사용하여 모델을 학습할 수 있습니다. 또한, MPS는 Apple Silicon (M1, M2 등)에서 GPU를 사용할 수 있는 백엔드입니다. 아래 코드를 통해 사용할 장치를 자동으로 설정할 수 있습니다:

import torch

def get_device():
    """
    PyTorch에서 사용할 수 있는 장치(CUDA, MPS, CPU)를 반환하는 함수.
    CUDA(GPU)가 사용 가능하면 CUDA, 
    Apple Silicon에서는 MPS, 
    그 외에는 CPU를 사용.
    
    Returns:
        device (str): 사용할 장치("cuda", "mps", "cpu").
    """
    device = (
        "cuda"  # GPU가 사용 가능하면 CUDA를 사용
        if torch.cuda.is_available()
        else "mps"  # Apple Silicon에서 GPU가 사용 가능하면 MPS 사용
        if torch.backends.mps.is_available()
        else "cpu"  # 그 외에는 CPU 사용
    )
    print(f"Using {device} device")
    return device

5. 모델 학습 (Training the Model)

모델을 학습시키는 단계입니다. 데이터셋을 여러 번 반복(에포크)하여 순전파와 역전파를 통해 모델의 매개변수를 업데이트합니다.

num_epochs = 10
device = get_device()  # 장치 설정

for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        # 입력 데이터와 레이블을 선택된 장치로 이동
        inputs, labels = inputs.to(device), labels.to(device)

        # 순전파
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 역전파 및 옵티마이저 단계
        optimizer.zero_grad()  # 기울기 초기화
        loss.backward()        # 역전파 계산
        optimizer.step()       # 가중치 업데이트

    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}')

6. 모델 평가 (Evaluating the Model)

학습된 모델을 평가하는 단계입니다. 보통 테스트 데이터셋을 사용하여 모델의 성능을 확인합니다. 평가 시에는 기울기 계산이 필요 없으므로 torch.no_grad()로 연산을 비활성화합니다.

model.eval()  # 평가 모드로 전환
correct = 0
total = 0

with torch.no_grad():  # 평가 시에는 기울기 계산을 하지 않음
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total:.2f}%')

7. 모델 저장 및 불러오기 (Saving & Loading the Model)

학습된 모델을 저장하고 나중에 불러올 수 있습니다.

# 모델 저장
torch.save(model.state_dict(), 'model.pth')

# 모델 불러오기
model = SimpleNN()
model.load_state_dict(torch.load('model.pth'))

추가 설명:

import os
import pandas as pd
from torchvision.io import read_image
from torch.utils.data import Dataset

class CustomImageDataset(Dataset):
    def __init__(self, img_dir, annotation_file, transform=None):
        self.img_dir = img_dir
        self.img_labels = pd.read_csv(annotation_file)
        self.transform = transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        return {"image": image, "label": label}

전체 코드 템플릿

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import os
import pandas as pd
from torchvision.io import read_image


# 1. 장치 설정 (Device Setting)
def get_device():
    """
    사용할 장치(CUDA, MPS, CPU)를 반환하는 함수.
    """
    device = (
        "cuda" if torch.cuda.is_available() else
        "mps" if torch.backends.mps.is_available() else
        "cpu"
    )
    print(f"Using {device} device")
    return device


# 2. 사용자 정의 데이터셋 클래스
class CustomImageDataset(Dataset):
    def __init__(self, img_dir, annotation_file, transform=None):
        """
        데이터셋 초기화 함수.
        """
        self.img_dir = img_dir
        self.img_labels = pd.read_csv(annotation_file)
        self.transform = transform

    def __len__(self):
        """
        데이터셋의 샘플 수를 반환하는 함수.
        """
        return len(self.img_labels)

    def __getitem__(self, idx):
        """
        주어진 인덱스에 해당하는 샘플을 반환하는 함수.
        """
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        return {"image": image, "label": label}


# 3. 데이터 로더 생성 함수
def create_dataloader(dataset, batch_size=32, shuffle=True):
    """
    DataLoader를 생성하는 함수.
    """
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)


# 4. 모델 정의 (Model Definition)
class SimpleNN(nn.Module):
    def __init__(self, input_size=784, hidden_size=128, output_size=10):
        """
        신경망 모델 초기화 함수.
        """
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        """
        순전파 함수.
        """
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 5. 손실 함수 및 옵티마이저 설정 함수
def setup_optimizer_and_criterion(model, lr=0.001):
    """
    손실 함수와 옵티마이저를 설정하는 함수.
    """
    criterion = nn.CrossEntropyLoss()  # 교차 엔트로피 손실 함수
    optimizer = optim.Adam(model.parameters(), lr=lr)  # Adam 옵티마이저
    return criterion, optimizer


# 6. 모델 학습 함수
def train_model(model, dataloader, criterion, optimizer, num_epochs=10, device="cpu"):
    """
    모델을 학습시키는 함수.
    """
    model.to(device)

    for epoch in range(num_epochs):
        model.train()  # 학습 모드로 전환
        running_loss = 0.0

        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)

            # 순전파
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # 역전파 및 옵티마이저 단계
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(dataloader):.4f}')


# 7. 모델 평가 함수
def evaluate_model(model, test_loader, device="cpu"):
    """
    학습된 모델을 평가하는 함수.
    """
    model.eval()  # 평가 모드로 전환
    correct = 0
    total = 0

    with torch.no_grad():  # 평가 시에는 기울기 계산을 하지 않음
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Accuracy: {accuracy:.2f}%')
    return accuracy


# 8. 모델 저장 함수
def save_model(model, file_path='model.pth'):
    """
    학습된 모델을 저장하는 함수.
    """
    torch.save(model.state_dict(), file_path)
    print(f'Model saved to {file_path}')


# 9. 모델 불러오기 함수
def load_model(file_path, input_size=784, hidden_size=128, output_size=10):
    """
    저장된 모델을 불러오는 함수.
    """
    model = SimpleNN(input_size, hidden_size, output_size)
    model.load_state_dict(torch.load(file_path))
    print(f'Model loaded from {file_path}')
    return model


# 전체 워크플로우
def run_training_pipeline(img_dir, annotation_file, test_loader, num_epochs=10, batch_size=32, lr=0.001):
    """
    전체 모델 학습, 평가, 저장 과정을 실행하는 함수.
    """
    # 1. 장치 설정
    device = get_device()

    # 2. 데이터셋 및 데이터로더 생성
    dataset = CustomImageDataset(img_dir=img_dir, annotation_file=annotation_file)
    dataloader = create_dataloader(dataset, batch_size=batch_size)

    # 3. 모델 정의
    model = SimpleNN()

    # 4. 손실 함수 및 옵티마이저 설정
    criterion, optimizer = setup_optimizer_and_criterion(model, lr=lr)

    # 5. 모델 학습
    train_model(model, dataloader, criterion, optimizer, num_epochs=num_epochs, device=device)

    # 6. 모델 평가
    evaluate_model(model, test_loader, device=device)

    # 7. 모델 저장
    save_model(model, 'model.pth')

세부 설명

  1. 장치 설정 (get_device): 사용 가능한 장치(CUDA, MPS, CPU)를 선택하는 함수를 정의하였습니다.

  2. 데이터셋 클래스 (CustomImageDataset): 이미지와 라벨을 로드하는 커스텀 데이터셋 클래스입니다.

  3. 데이터 로더 생성 함수 (create_dataloader): DataLoader를 통해 배치 단위로 데이터를 불러옵니다.

  4. 모델 정의 (SimpleNN): 간단한 완전 연결 신경망(FCN)을 정의하는 클래스입니다.

  5. 손실 함수 및 옵티마이저 설정 함수 (setup_optimizer_and_criterion): 손실 함수와 옵티마이저를 설정하는 함수입니다.

  6. 모델 학습 함수 (train_model): 입력된 데이터로 모델을 학습시키는 함수입니다.

  7. 모델 평가 함수 (evaluate_model): 학습된 모델을 평가하는 함수입니다.

  8. 모델 저장 함수 (save_model): 학습된 모델을 저장하는 함수입니다.

  9. 모델 불러오기 함수 (load_model): 저장된 모델을 불러오는 함수입니다.


참고자료

https://velog.io/@ohado/PyTorch-%ED%8C%8C%EC%9D%B4%ED%86%A0%EC%B9%98-%EC%A0%84%EC%B2%B4-%EA%B5%AC%EC%A1%B0

profile
AI Engineer / 의료인공지능

0개의 댓글