【1】人工智慧訓練 - 選擇最優模型

import cv2 as cv
from sklearn.model_selection import train_test_split
import torch
from torch.utils.data import DataLoader, TensorDataset
from torchvision import models
from torch import nn, optim
import os
from PIL import Image

data_set = []       #資料集
# 匯入本地圖片
folder_path = 'sample/'  #相對路徑
data_set = []   #存放轉換後的圖片資料集
labels = []
for i in range(2):
    for filename in os.listdir(folder_path+str(i)):
        img = cv.imread(os.path.join(folder_path+str(i) + '/', filename))


        img  = cv.resize(img, (112, 112))
        img = img.reshape(3, 112, 112)
        data_set.append(img)
        labels.append(i)    #將0123放到標籤列表中
# print(data_set)
# print(labels)
# 提取特徵和目標變數
X = data_set    #訓練資料
y = [0] * 100 + [1] * 200  #資料圖片的數量

# 將資料集分為訓練集和測試集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 將特徵資料轉換為PyTorch張量
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
# X_train_tensor = torch.tensor(X_train)
# X_test_tensor = torch.tensor(X_test)

# 將目標變數轉換為PyTorch張量
y_train_tensor = torch.tensor(y_train, dtype=torch.long)  # 假設目標變數是整數類別
y_test_tensor = torch.tensor(y_test, dtype=torch.long)

# 建立訓練集和測試集的資料集物件
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)

# 建立訓練集和測試集的資料載入器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 選擇模型
model = models.resnet18(num_classes=2)  #殘差網路18個隱層
# 將模型載入到裝置上
model.to(device)
# 選擇損失函式
loss_fn = nn.CrossEntropyLoss()
# 選擇最佳化函式
opt = optim.SGD(model.parameters(), lr=0.3)

# 定義訓練
def train(epoch):
    model.train()
    i = 0
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        # 計算損失
        loss = loss_fn(outputs, labels)
        # 梯度初始化
        opt.zero_grad()
        loss.backward()
        opt.step()
        print("\rEpoch: {} [batch: {}/{} loss:{:.3f}]".format(epoch, i + 1, len(train_loader), loss), end=' ')
        i += 1
    print()
    return loss.item()

# 定義測試
def valid():
    model.eval()
    correct_qty = 0
    with torch.no_grad():
        for image, label in test_loader:
            image = image.to(device)
            label = label.to(device)
            output = model(image)
            pred = torch.argmax(output, dim=1)
            correct_qty += (pred == label).sum().item()

    acc = correct_qty / len(test_loader.dataset)
    print("準確率：", acc)
    return acc

acc_best = 0
for epoch in range(10):
    loss = train(epoch)
    acc = valid()
    if acc_best < acc:
        acc_best = acc
        print("儲存模型")
        torch.save(model.state_dict(), "model-4.pt")
    else:
        print("不儲存模型")