PyTorch基础-猫狗分类实战-10

训练模型并保存

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets,transforms,models
from torch.utils.data import Dataset
import sys

# 数据预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),# 对图像进行随机裁剪
    transforms.RandomRotation(20),# 随机旋转角度
    transforms.RandomHorizontalFlip(p=0.5),# 随机水平翻转
    transforms.ToTensor()# 变成tensor格式
]) # 数据增强

# 读取数据
root = "image"
train_dataset = datasets.ImageFolder(root + "/train",transform)
test_dataset = datasets.ImageFolder(root + "/test",transform)

# 导入数据
train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=8,shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset,batch_size=8,shuffle=True)

classes = train_dataset.classes
classes_index = train_dataset.class_to_idx
print(classes)
print(classes_index)

model = models.vgg16(pretrained=True)# 载入vgg16预训练模型
print(model)

for param in model.parameters():
    param.requires_grad = False
# 构建新的全连接层
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088,100),
                                       torch.nn.ReLU(),
                                       torch.nn.Dropout(p=0.5),
                                       torch.nn.Linear(100,2))

LR = 0.0003
# 定义代价函数
entropy_loss = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(model.parameters(),LR)

def train():
    model.train()
    for i,data in enumerate(train_loader):
        # 获得数据和对应的标签
        inputs,labels = data
        # 获得模型预测结果（64，10）
        out = model(inputs)
        # 交叉熵代价函数out（batch.C),labels(batch)
        loss = entropy_loss(out,labels)
        # 梯度清零
        optimizer.zero_grad()
        # 计算梯度
        loss.backward()
        # 修改权值
        optimizer.step()
        
def test():
    model.eval()
    correct = 0
    for i,data in enumerate(test_loader):
        # 获得数据和对应的标签
        inputs,labels = data
        # 获得模型预测结果
        out = model(inputs)
        # 获得最大值，以及最大值所在的位置
        _,predicted = torch.max(out,1)
        # 预测正确的数量
        correct += (predicted == labels).sum()
    print("test acc:{0}".format(correct.item()/len(test_dataset)))
    
    correct = 0
    for i,data in enumerate(train_loader):
        # 获得数据和对应的标签
        inputs,labels = data
        # 获得模型预测结果
        out = model(inputs)
        # 获得最大值，以及最大值所在的位置
        _,predicted = torch.max(out,1)
        # 预测正确的数量
        correct += (predicted == labels).sum()
    print("train acc:{0}".format(correct.item()/len(train_dataset)))
        

for  epoch in range(5):
    print("epoch:",epoch)
    train()
    test()
torch.save(model.state_dict(),"cat_dog.pth") # 保存模型

加载模型进行预测

import torch
import numpy as np
from PIL import Image
from torchvision import transforms,models

model = models.vgg16(pretrained=True)
# 构建新的全连接层
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088,100),
                                       torch.nn.ReLU(),
                                       torch.nn.Dropout(p=0.5),
                                       torch.nn.Linear(100,2))
model.load_state_dict(torch.load("cat_dog.pth")) # 加载模型
model.eval() # 预测模式

label = np.array(["cat","dog"])
# 数据预处理
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor()
])

# 预测函数
def predict(image_path):
    # 打开图片
    img = Image.open(image_path)
    # 数据处理，增加一个维度
    img = transform(img).unsqueeze(0)
    # 预测得到的结果
    outputs = model(img)
    # 获得最大值所在位置
    _,predicted  = torch.max(outputs,1)
    # 转换为类别名称
    print(label[predicted.item()])

predict("image/test/cat/cat.1490.jpg")