【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("不保存模型")