【PyTorch】图像多分类项目

【PyTorch】图像多分类项目部署

如果需要在独立于训练脚本的新脚本中部署模型，这种情况模型和权重在内存中不存在，因此需要构造一个模型类的对象，然后将存储的权重加载到模型中。

加载模型参数，验证模型的性能，并在测试数据集上部署模型

from torch import nn
from torchvision import models

# 定义一个resnet18模型，不使用预训练参数
model_resnet18 = models.resnet18(pretrained=False)
# 获取模型的全连接层的输入特征数
num_ftrs = model_resnet18.fc.in_features
# 定义分类的类别数
num_classes=10
# 将全连接层的输出特征数改为分类的类别数
model_resnet18.fc = nn.Linear(num_ftrs, num_classes)

import torch 
path2weights="./models/resnet18_pretrained.pt"
# 加载预训练的ResNet18模型权重
model_resnet18.load_state_dict(torch.load(path2weights))
# 将ResNet-18模型设置为评估模式
model_resnet18.eval();
# 检查CUDA是否可用
if torch.cuda.is_available():
    # 如果可用，将设备设置为CUDA
    device = torch.device("cuda")
    # 将模型移动到CUDA设备上
    model_resnet18=model_resnet18.to(device)

def deploy_model(model,dataset,device, num_classes=10,sanity_check=False):

    # 获取数据集的长度
    len_data=len(dataset)
    
    # 初始化输出张量
    y_out=torch.zeros(len_data,num_classes)
    
    # 初始化真实标签张量
    y_gt=np.zeros((len_data),dtype="uint8")
    
    # 将模型移动到指定设备
    model=model.to(device)
    
    # 初始化时间列表
    elapsed_times=[]
    with torch.no_grad():
        for i in range(len_data):
            # 获取数据集中的一个样本
            x,y=dataset[i]
            # 将真实标签存入张量
            y_gt[i]=y
            # 记录开始时间
            start=time.time()    
            # 将输入数据传入模型进行预测
            yy=model(x.unsqueeze(0).to(device))
            # 将预测结果存入张量
            y_out[i]=torch.softmax(yy,dim=1)
            # 计算预测时间
            elapsed=time.time()-start
            # 将预测时间存入列表
            elapsed_times.append(elapsed)

            # 如果进行完整性检查，则跳出循环
            if sanity_check is True:
                break

    # 计算平均预测时间
    inference_time=np.mean(elapsed_times)*1000
    # 打印平均预测时间
    print("average inference time per image on %s: %.2f ms " %(device,inference_time))
    # 返回预测结果和真实标签
    return y_out.numpy(),y_gt

from torchvision import datasets
import torchvision.transforms as transforms

# 数据转换
data_transformer = transforms.Compose([transforms.ToTensor()])

path2data="./data"

# 加载数据
test0_ds=datasets.STL10(path2data, split='test', download=True,transform=data_transformer)
print(test0_ds.data.shape)

from sklearn.model_selection import StratifiedShuffleSplit

# 创建StratifiedShuffleSplit对象，设置分割次数为1，测试集大小为0.2，随机种子为0
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=0)

# 获取test0_ds的索引
indices=list(range(len(test0_ds)))

# 获取test0_ds的标签
y_test0=[y for _,y in test0_ds]

# 对索引和标签进行分割
for test_index, val_index in sss.split(indices, y_test0):
    # 打印测试集和验证集的索引
    print("test:", test_index, "val:", val_index)
    # 打印测试集和验证集的大小
    print(len(val_index),len(test_index))

from torch.utils.data import Subset

# 从test0_ds中选取val_index索引的子集，赋值给val_ds
val_ds=Subset(test0_ds,val_index)
# 从test0_ds中选取test_index索引的子集，赋值给test_ds
test_ds=Subset(test0_ds,test_index)
# 定义均值
mean=[0.4467106, 0.43980986, 0.40664646]
# 定义标准差
std=[0.22414584,0.22148906,0.22389975]
# 定义一个名为test0_transformer的变量，用于将一系列的图像变换操作组合在一起
test0_transformer = transforms.Compose([
    # 将图像转换为Tensor类型
    transforms.ToTensor(),
    # 对图像进行归一化操作，使用mean和std作为均值和标准差
    transforms.Normalize(mean, std),
    ])   
# 将test0_transformer赋值给test0_ds的transform属性
test0_ds.transform=test0_transformer

import time
import numpy as np

# 调用deploy_model函数，传入model_resnet18,val_ds,device和sanity_check参数，返回y_out和y_gt
y_out,y_gt=deploy_model(model_resnet18,val_ds,device=device,sanity_check=False)
# 打印y_out和y_gt的形状
print(y_out.shape,y_gt.shape)

from sklearn.metrics import accuracy_score

# 将y_out中的最大值索引赋值给y_pred
y_pred = np.argmax(y_out,axis=1)
# 打印y_pred和y_gt的形状
print(y_pred.shape,y_gt.shape)

# 计算并打印y_pred和y_gt的准确率
acc=accuracy_score(y_pred,y_gt)
print("accuracy: %.2f" %acc)

 

# 部署模型，得到预测结果和真实标签
y_out,y_gt=deploy_model(model_resnet18,test_ds,device=device)

# 取出预测结果中概率最大的类别
y_pred = np.argmax(y_out,axis=1)

# 计算准确率
acc=accuracy_score(y_pred,y_gt)

# 打印准确率
print(acc)

from torchvision import utils
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline
np.random.seed(1)

# 定义一个函数，用于显示图像
def imshow(inp, title=None):
    # 定义图像的均值和标准差
    mean=[0.4467106, 0.43980986, 0.40664646]
    std=[0.22414584,0.22148906,0.22389975]
    # 将图像从tensor转换为numpy数组，并转置
    inp = inp.numpy().transpose((1, 2, 0))
    # 将均值和标准差转换为numpy数组
    mean = np.array(mean)
    std = np.array(std)
    # 将图像的像素值进行归一化
    inp = std * inp + mean
    # 将像素值限制在0和1之间
    inp = np.clip(inp, 0, 1)
    # 显示图像
    plt.imshow(inp)
    # 如果有标题，则显示标题
    if title is not None:
        plt.title(title)
    # 暂停0.001秒
    plt.pause(0.001) 

# 定义网格大小
grid_size=16
# 随机生成4个索引
rnd_inds=np.random.randint(1,len(test_ds),grid_size)
# 打印随机生成的索引
print("image indices:",rnd_inds)

# 根据索引获取对应的图像和标签
x_grid_test=[test_ds[i][0] for i in rnd_inds]
y_grid_test=[(y_pred[i],y_gt[i]) for i in rnd_inds]

# 将图像转换为网格
x_grid_test=utils.make_grid(x_grid_test, nrow=4, padding=2)
# 打印网格的形状
print(x_grid_test.shape)

# 设置图像的大小
plt.rcParams['figure.figsize'] = (10, 10)
# 显示网格
imshow(x_grid_test,y_grid_test)