4-2.py

import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from torch import nn,optim
import os

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

transform = transforms.ToTensor()
train_DS = datasets.MNIST(root ='./data',train=True,download=True,transform=transform)
test_DS = datasets.MNIST(root='./data',train=False,download=True,transform=transform)
#root로 어디에 데이터 저장할 지 표시 train으로 train 데이터 셋인지 아닌 지 download로 data 폴더 안에 없을 때 다운 할지 안할지 transform으로 어떤 형태로 변환할건지

#train_DS.classes라벨이 뭐가 있는 지 알려줌

#train_DS.class_to_idx
#오른쪽에 있는게 레이블 왼쪽은 클래스명

#print(train_DS.data.shape) train_DS의 shape를 보여줌
'''
plt.imshow(train_DS.data[0], cmap="gray")
plt.colorbar()
plt.show()
'''

BATCH_SIZE = 32#보통 2의 제곱수으로
train_DL = torch.utils.data.DataLoader(train_DS,batch_size=BATCH_SIZE,shuffle=True)
test_DL = torch.utils.data.DataLoader(train_DS,batch_size=BATCH_SIZE,shuffle=True)
#train_DS에 대해서 batch_size는 하이퍼파라미터로 가져오고,shuffle=True로 데이터를 섞어서 무작위로 가지고 오게 함

#train_DL.dataset DL은 데이터 셋도 들고 있음

x_batch, y_batch = next(iter(train_DL))#데이터셋에서 한 배치 사이즈만큼만 꺼내기
#ToTensor의 역할
#1. Tensor로 바꾸고
#2. 개체 행렬로 바꾸고
#3. 0~1로 바꿈(int -> float)
#print(x_batch.size)해보면 32,1,28,28이 나옴 32가 개 1,28,28이 채 흑백이니까 1로 RGB면 3임

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Sequential(nn.Linear(28*28,100),
                                    nn.ReLU(),
                                    nn.Linear(100,10))
    def forward(self,x):
        x = torch.flatten(x, start_dim=1)#개체 행렬인거를 개x채로 만들어주기 위해서 쫙쫙 펴주는 함수
        x = self.linear(x)
        return x
model = MLP().to(DEVICE)

LR = 1e-3

EPOCH = 5

Loss_fn = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr = LR)


def Train(EPOCH,model,optimizer,Loss_fn,train_DL):
    model.train()
    loss_history = []
    NoT = len(train_DL.dataset)
    for ep in range(EPOCH):
        rloss = 0
        for x_batch, y_batch in train_DL:#배치 사이즈(32)만큼 한 for문에서 갖고오는 걸 train_DL의 크기 만큼 반복
            x_batch = x_batch.to(DEVICE)
            y_batch = y_batch.to(DEVICE)
            y_hat = model(x_batch)
            loss = Loss_fn(y_hat,y_batch)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss_b = loss.item() * x_batch.shape[0]
            #batch_size만큼 가져오다보면 50에서 32 가져오고 나머지가 18인데 그러면 어쩔 수 없이 18을 가져오게 되지만 여기서 시그마로 1부터 32까지 더하던게 1부터 18까지 더하는 걸로 바뀌고
            #평균 값도 1/32 하던걸 1/18로 하니까 둘을 더한 평균이 제대로 나오지 않을거임 그래서 여기서 1/32같이 하는 거 에서 batch_size 만큼 즉 32를 곱해버려서 나중에 그 2개를 더한 뒤 평균을 내겠다는 전략
            rloss += loss_b
        loss_e = rloss/NoT
        loss_history += [loss_e]
        print(f"Epoch: {ep+1} Loss {loss_e:.3f}")
        print("-"*20)

    return loss_history
loss_history = Train(EPOCH,model,optimizer,Loss_fn,train_DL)

save_model_path = "./model/model.pth"

#model.state_dict() 모델의 파라미터들을 딕셔너리로 표현해줌

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

load_model = MLP().to(DEVICE)

load_model.load_state_dict(torch.load(save_model_path, map_location=DEVICE))#map_location 있어야 GPU로 학습 했던 거 device로 불러올 수 있다


def Test(model,Test_DL):
    model.eval()
    with torch.no_grad():
        rcorrect = 0
        for x_batch, y_batch in Test_DL:
            x_batch = x_batch.to(DEVICE)
            y_batch = y_batch.to(DEVICE)
            y_hat = model(x_batch)
            pred = y_hat.argmax(dim=1)
            #32,10이 출력으로 나올테고 32개의 행에 각각 어떤 숫자로 예측 했는 지가 열에 담겨 있을테니까 argmax로 가장 높은 값 즉 가장 출력이 높은 인덱스를 가져와서 어떤 값으로 예측한 지 보는 거임
            rcorrect_a = torch.sum(pred == y_batch)
            rcorrect += rcorrect_a
        correct = rcorrect/len(Test_DL.dataset) * 100
    print(f"test_correct: {rcorrect/len(Test_DL.dataset)} ({correct:.1f}%)")
    return correct
Test(load_model,test_DL)