파이썬 AI [입문] Pytorch : 파이토치(2) 기본모델 구축 및 저장하기 (CNN)

 

오늘도 역시 파이토치 배워봅시다

파이토치에서 기본모델을 구축하고 저장하는 방법 및 사용법 알아보죠 !

정리 먼저 읽어주시고, 그다음 코드를 따라가면서 주석부분을 보시면서 공부하면 됩니다 ! : )

 

정리

1. 데이터 전처리 해주기

2. 모델 만들기

3. 훈련데이터로 모델 사용하기

4. 훈련된 모델 저장하기

5. 저장된 모델을 불러와서 예측하기

출처 입력

코드

import torch
import torchvision
import torchvision.transforms as transforms # 데이터를 불러오면서 데이터 전처리 하기위한 라이브러리
from torch.utils.data import Dataset # Dataset으 우리가 튜닝을 할 때 사용
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim #그레디언트를 구할 수 있는 optimizer


### 기본 모델 구축 및 저장하기

# Load Data set
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # 데이터 전처리를 위한 작업

trainset = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) # 훈련데이터 설정
trainloader = torch.utils.data.DataLoader(trainset, batch_size=8, shuffle=True,num_workers=2) # 모델 사용하기 위한 배치사이즈 설정

testset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) # 테스트데이터 설정
testloader = torch.utils.data.DataLoader(testset, batch_size=8, shuffle=True, num_workers=2) # 모델 사용하기 위한 배치사이즈 설정

classes = ('plane', 'car', 'bird', 'cat', 'deer','dog','frog','horse','ship','truck')

# Build a Model
class Net(nn.Module): #nn.Module을 상속 받는다
    def __init__(self): # 연산을 세팅하는 부분(딥러닝에서 사용할 연산들을 세팅하는 부분)
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3,6,5)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6,16,5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self,x): # 세팅된 연산의 순서를 정해주는 부분
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

print(net)

# Implement the model with training set
criterion = nn.CrossEntropyLoss() # 분류 문제에 쓰이는 loss 함수
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # 정의해준 모델을 넣어주며, learning rate은 0.001,

for epoch in range(3): # loop over the data multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader,0):
        inputs, labels = data # get the inputs; data is a list of [inputs,labels]

        optimizer.zero_grad() # zero the parameter gradient

        outputs = net(inputs) # forward + backward + optimize
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.itme() # print statics
        if i % 2000 == 1999:
            print('[%d, %5d] loss : %.3f' %(epoch+1,i+1,running_loss/2000))
            running_loss = 0.0
print('Finish Training')

# Save the trained Model
path = "./cifar_net.pth"
torch.save(net.state_dict(),path)
# Load the pre-trained Model
net = Net()
net.load_state_dict(torch.load(path))

# Predict Model
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted =  torch.max(outputs.data,1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print("Accuracy of the network on the 10000 test images : %d"%(100*correct/total))