AI 개발을 하면서 성능개선과 유지보수를 잘 하려면 파일 구조를 정형화하고 구조화 시켜야 한다.
이 글에서는 내가 사용하는 파일 구조와 이 파일 구조의 장점을 알려주겠다.
내 파일 구조를 먼저 보여주겠다
나는 AI를 학습 시킬 때 주로 4개의 파일로 나눠 코드를 짠다.
파일 하나하나씩 그 역할을 설명하고 예시 코드를 첨부하겠다.
예시 코드는 간단한 MNIST 분류 모델이다.
I . dataloader.py
이 파일은 데이터를 다운로드하거나 불러오고 데이터 전처리와 데이터 셋 분리 및 PyTorch의 dataloader 객체로 변환하는 역할을 한다.
예시 코드)
from torchvision.datasets import MNIST
from torchvision.transforms import v2
from torch.utils.data import DataLoader
from torch.utils.data import random_split
import torch
def get_transform():
transform = v2.Compose([
v2.RandomAffine(
degrees=5,
translate=(0.1, 0.1)
), # ±5도 기울이기, ±10% 움직이기
v2.ToImage(), # PIL -> tv_tensor.Image
v2.ToDtype(dtype=torch.float32, scale=True) # dtype을 float로 변환 및 0~1의 사이값으로 scale 해줌
])
return transform
def get_dataset(cfg, transform):
train_dataset = MNIST(root="./data", train=True, transform=transform, download=True) # get train dataset
train_dataset, valid_dataset = random_split(train_dataset, [int(len(train_dataset) * cfg.data.train_ratio), len(train_dataset)-int(len(train_dataset) * cfg.data.train_ratio)]) # split train, valid dataset
test_dataset = MNIST(root="./data", train=False, transform=transform, download=True) # get test dataset
return train_dataset, valid_dataset, test_dataset
def get_loaders(cfg, train_dataset, valid_dataset, test_dataset):
train_loader = DataLoader(train_dataset, batch_size=cfg.data.batch_size, shuffle=True) # get train dataloader
valid_loader = DataLoader(valid_dataset, batch_size=cfg.data.batch_size, shuffle=True) # get valid dataloader
test_loader = DataLoader(test_dataset, batch_size=cfg.data.batch_size, shuffle=False) # get test dataloader
return train_loader, valid_loader, test_loader
II . model.py
이 파일은 모델을 선언하는 파일이다.
예시 코드)
import torch.nn as nn
class SimpelModel(nn.Module):
def __init__(self):
super(SimpelModel, self).__init__()
self.fc1 = nn.Linear(28*28, 10) # input_size = 28 * 28, output_size = 10
def forward(self, x):
x = x.view(x.size(0), -1) # flatten
x = self.fc1(x) # feedforward
return x
III . trainer.py
이 파일은 train 과정과 validation 과정을 작성한 파일이다.
나는 tensorboard를 사용하기 때문에 writer를 작동하는 코드가 끼어있다. 만약 tensorboard를 사용하지 않는다면 writer 코드 부분을 삭제하면 된다.
예시 코드)
import torch
import logging
def train(model, train_loader, valid_loader, criterion, optimizer, epochs, writer):
best_accuracy = 0.0 # best accuracy 초기화
best_model_weights = None # best model weights 초기화
model.train() # batch_norm, drop_out 활성화
for epoch in range(epochs):
running_loss = 0.0 # loss 초기화
for images, labels in train_loader:
outputs = model(images) # images feedforward
loss = criterion(outputs, labels) # loss 계산
optimizer.zero_grad() # gradient 초기화
loss.backward() # backpropagation
optimizer.step() # step
running_loss += loss.item() # sum of loss
avg_train_loss = running_loss / len(train_loader) # loss avg
writer.add_scalar('Loss/train', avg_train_loss, epoch) # tensorboard에 저장
model.eval() # batch_norm, drop_out 비활성화
correct, total = 0, 0 # correct, total 초기화
with torch.no_grad(): # 역전파와 gradient 계산 비활성화
for images, labels in valid_loader:
outputs = model(images) # images feedforward
_, predicted = torch.max(outputs, 1) # 확률 값이 가장 큰 index 값 저장
total += labels.size(0) # 데이터의 개수 저장
correct += (predicted == labels).sum().item() # 맞은 데이터 개수 저장
accuracy = 100 * correct / total # 정확도 계산
writer.add_scalar('Accuracy/validation', accuracy, epoch) # epoch에 따른 accracy 저장
if accuracy > best_accuracy: # best accuracy 보다 현 epoch의 accuracy가 높다면
best_accuracy = accuracy # best accuracy로 저장
best_model_weights = model.state_dict().copy() # best accuracy의 model weights 저장
logging.info(f"Epoch [{epoch+1}/{epochs}], Loss: {avg_train_loss:.4f}, Validation Accuracy: {accuracy:.2f}%")
model.train()
if best_model_weights:
model.load_state_dict(best_model_weights) # model에 best accuracy의 weight를 가져옴
torch.save(best_model_weights, 'best_model.pth') # model 저장
logging.info(f"Best model weights saved with accuracy: {best_accuracy:.2f}")
def evaluate(model, test_loader, writer):
model.eval()
correct, total = 0, 0
with torch.no_grad():
for images, labels in test_loader:
outputs = model(images) # images feedforward
_, predicted = torch.max(outputs, 1) # 확률값이 가장큰 index 저장
total += labels.size(0) # 데이터의 개수 저장
correct += (predicted == labels).sum().item() # 맞은 데이터 개수 저장
accuracy = 100 * correct / total # 정확도 계산
logging.info(f"Accuracy: {accuracy:.2f}%")
writer.add_scalar('Accuracy/test', accuracy) # accracy 저장
writer.close() # tensorboard writer close
IV . train.py
이 파일은 위에 있는 파일들을 가지고 진짜 모델을 train 시키는 파일이다.
나는 omegaconf와 hydra를 사용해서 config 파일이 하나 더 있긴한데 사용하지 않아도 된다.
예시 코드)
import hydra
from omegaconf import OmegaConf
from torch.utils.tensorboard import SummaryWriter
import torch.optim as optim
from dataloader import get_transform, get_dataset, get_loaders
from model import SimpelModel
import logging
import torch.nn as nn
from trainer import train, evaluate
@hydra.main(version_base=None, config_path="./config", config_name='train')
def main(cfg):
OmegaConf.to_yaml(cfg) # get config
transform = get_transform() # get transform
train_dataset, valid_dataset, test_dataset = get_dataset(cfg, transform) # get dataset
logging.info(f"train_dataset: {len(train_dataset)}, valid_dataset: {len(valid_dataset)}, test_dataset: {len(test_dataset)}")
train_loader, valid_loader, test_loader = get_loaders(cfg, train_dataset, valid_dataset, test_dataset) # get dataloader
logging.info(f"train_loader: {len(train_loader)}, valid_loader: {len(valid_loader)}, test_loader: {len(test_loader)}")
model = SimpelModel() # get model
criterion = nn.CrossEntropyLoss() # get loss function
optimizer = optim.SGD(model.parameters(), lr=cfg.train.lr) # get optimizer
writer = SummaryWriter() # get writer
train(model, train_loader, valid_loader, criterion, optimizer, cfg.train.epochs, writer) # train
evaluate(model, test_loader, writer) # evaluate
if __name__ == '__main__':
main()
V . 이렇게 하면 좋은 점?
파일을 이렇게 구조화하면 모델을 갑자기 바꿔야할 때 바로 model.py 파일로 가면 되고 데이터를 바꾸고 싶을 때도 dataloader.py로 가면 되기 때문에 정말 손쉽게 교체할 수 있다. 그리고 이런 파일 구조로 코드를 짜면 협업할 때도 어떤 코드가 무슨 역할을 하는지 잘 구분되어 편리하다.