基于卷积神经网络的图片分类

pytorch中文教程

项目地址

1动物分类__cnn

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

# 数据预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),# 对图像进行随机的crop以后再resize成固定大小
    transforms.RandomRotation(20), # 随机旋转角度
    transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转
    transforms.ToTensor()  #转换PIL图像为张量
])
 
# 读取数据
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)
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))
print(model)

LR = 0.0001
# 定义代价函数 交叉熵损失函数
entropy_loss = nn.CrossEntropyLoss()
# 定义优化器 随机梯度下降
optimizer = optim.SGD(model.parameters(), LR, momentum=0.9)


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)
        # 梯度清0
        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(0, 10):
    print('epoch:',epoch)
    train()
    test()
    
torch.save(model.state_dict(), 'cat_dog_cnn.pth')

2测试程序

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_cnn.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.1000.jpg')