predict.py

import torch
from model import MLP
from dataloader import get_dataset

_,Test_DL = get_dataset(32)

model = MLP()

model.load_state_dict(torch.load("./model/bs_32_ep_5_lr_1e-3_lf_CE_optim_Adam.pth"))

def Test(model,Test_DL):
    model.eval()
    with torch.no_grad():
        rcorrect = 0
        for x_batch, y_batch in Test_DL:
            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(model,Test_DL)