[PyTorch] 패션 MNIST

패션 MNIST

 

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import os

# 트레이닝 데이터셋 다운로드
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

# 테스트 데이터셋 다운로드
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

img = test_data[0][0].numpy() * 255
img = img.astype(np.int32).reshape(28,28)
os.environ['KMP_DUPLICATE_LIB_OK']='True'
plt.imshow(img, cmap="gray")

 

데이터 로더

 

batch_size = 64

train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

 

GPU 설정

 

device = "cuda" if torch.cuda.is_available() else  "cpu"

# 또는
'''
if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
'''

print(f"Using {device} device")

 

신경망

 

class NerualNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
        
    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NerualNetwork().to(device) # 변수를 cuda로 보냄
print(model)

NerualNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

 

하이퍼파라미터

 

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) # 0.001

 

학습

 

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train() # 모델이 학습하도록 트레인 단계 활성화
    
    for batch, (X,y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        
        # 모델의 에러를 계산
        pred = model(X) # 포워드 연산
        loss = loss_fn(pred, y) # loss 연산
        
        # 역전파
        optimizer.zero_grad() # 매 이터레이션마다 그래디언트 초기화
        loss.backward() # 역전파가 진행되어 w,r,t 그래프잎의 기울기 계산
        optimizer.step() # 가중치 업데이트
        
        if batch % 100 == 0: # 매 100회 배치를 돌 때마다 리포트 수행
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
            
def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval() # 평가 수행
    test_loss, correct = 0, 0
    
    with torch.no_grad(): # 그라데이션 계산 비활성화
        for X,y in dataloader:
            X, y = X.to(device), y.to(device)

            # 모델의 에러 계산
            pred = model(X) # 포워드 연산
            test_loss += loss_fn(pred, y).item() # loss 연산
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f}\n")

    
epochs = 5

for t in range(epochs):
    print(f'Epoch {t+1}\n--------------------------------------')
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
    
print('Done.')

 

모델 저장 / 불러오기

 

torch.save(model.state_dict(), "model.pth")
print("Pytorch 모델의 상태를 model.pth에 저장")

mode = NerualNetwork()
model.load_state_dict(torch.load("model.pth"))

classes = [
    "T-shirt/top",
    "Trouse",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval()
x, y = test_data[0][0], test_data[0][1]

with torch.no_grad():
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

Predicted: "Ankle boot", Actual: "Ankle boot"