import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from config import args
from MyDataset import MyDataset
from model import Cnn1, Cnn2
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

# 获取测试集数据
transform = transforms.ToTensor()
test_set = MyDataset(args.root + '/HWDB1.1tst_gnt', transforms=transform)
test_loader = DataLoader(test_set, batch_size=args.batch_size)

# 加载 Train 模型
device = torch.device('cuda' if args.cuda else 'cpu')
model = eval(args.model + '().to(device)')  # Cnn().to(device)
model.load_state_dict(torch.load(args.root + '/param', map_location=device))
model.eval()
criterion = nn.CrossEntropyLoss()
eval_acc = 0
eval_loss = 0

# 测试
for data in test_loader:
    image, label = data
    image, label = image.to(device), label.to(device)

    out = model(image)  # 前行算法
    loss = criterion(out, label)  # 计算loss
    eval_loss += loss.item() * label.size(0)  # 计算总的loss

    _, pred = torch.max(out, 1)  # 预测结果
    num_correct = (pred == label).sum()  # 正确结果
    eval_acc += num_correct.item()  # 正确结果总数

    print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(test_set)), eval_acc / (len(test_set))))