# 평가 함수 def evaluate(model, test_loader, device): model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: images = images.to(device) labels = labels.to(device)
outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() ### 다시 train 모드로 model.train() return correct / total
train.py
# 패키지 임포트 import torch import torch.nn as nn from torch.optim import Adam
from utils.tools import evaluate from utils.getModules import load_dataloaders from networks.vgg import myVGG
# Hparam 설정 batch_size = 100 num_classes = 10 num_epochs = 10 # device 설정 device='cuda' if torch.cuda.is_available() else 'cpu'
## 데이터 처리하는 과정 (CIFAR10 활용) train_loader, test_loader = load_dataloaders(batch_size)
## 모델 만들기 ##
# 모델 생성에 사용이 될 Hparam 설정 ''' 상단에 전체 Hparam 적는 과정에 기입 ''' # 모델 객체를 생성 (설계도 + Hparam) model = myVGG(num_classes=num_classes).to(device)
best_acc = 0 # for loop를 돌면서 데이터를 불러오기 for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device)
# 불러온 데이터를 모델에 넣기 outputs = model(images) # 나온 출력물로 loss를 계산 loss = loss_func(outputs, labels)
# Loss로 back prop 진행 loss.backward() # optimizer를 이용해 최적화를 진행 optimizer.step() optimizer.zero_grad()
# 학습 중간에 if (i+1) % 100 == 0: print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item()}') # 평가를 진행해서 acc = evaluate(model, test_loader, device) # 성능이 좋으면 if acc > best_acc: best_acc = acc # 저장을 진행 torch.save(model.state_dict(), f'./best_model.ckpt')
infer.py
# 패키지를 임포트 import torch from PIL import Image from torch.nn.functional import softmax
from utils.getModules import load_transform from networks.vgg import myVGG
# device 설정 device='cuda' if torch.cuda.is_available() else 'cpu'
# 추론에 사용할 데이터를 준비 (전처리) example_image_path = 'plane.png' # 이미지 읽고 image = Image.open(example_image_path) # 크기 변환 (32 x 32) image = image.resize((32, 32))
# 학습에서 사용한 전처리 설정 transform = load_transform()
# 전처리 진행 image = transform(image) image = image.unsqueeze(0).to(device)
## 학습이 완료된 최고의 모델을 준비하기 num_classes = 10 model = myVGG(num_classes=num_classes)
# 속이 빈 모델에 학습된 모델의 wieght를 덮어 씌우기 weight = torch.load('best_model.ckpt', map_location=device) model.load_state_dict(weight) model = model.to(device)
# 준비된 데이터를 모델에 넣기 ### model 평가 모드 model.eval() output = model(image)
## 결과를 분석 # 추론 결과 만들어내기 output = output[0] probabilities = softmax(output, dim=0) # softmax : 확률값으로 변환 prob, index = torch.max(probabilities, dim=0) prob = prob * 100
