pytorch 图片分类，python 图片分类，net 图片分类

pytorch 图片分类，python 图片分类，net 图片分类，深度学习 图片分类

pytorch版本：1.5.0+cu101

全部源码，可以直接运行。

下载地址：https://download.csdn.net/download/TangLingBo/12598527

网络是用 net

如果出现需要下载的文件或者问题可以联系：QQ 1095788063

图片结构：

 

测试结果：

 

 

网络代码：

# encoding=utf-8
import torch.nn as nn
import torch.nn.functional as F


# 定义网络，继承torch.nn.Module
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)  # 卷积层
        self.pool = nn.MaxPool2d(2, 2)  # 池化层
        self.conv2 = nn.Conv2d(6, 16, 5)  # 卷积层
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)  # 6个输出

    def forward(self, x):  # 前向传播

        x = self.pool(F.relu(self.conv1(x)))  # F就是torch.nn.functional
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)  # .view( )是一个tensor的方法，使得tensor改变size但是元素的总数是不变的。
        #x = x.view(-1, 3 * 224 * 224)
        # 从卷基层到全连接层的维度转换

        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

 

训练代码：

import torch as t
import torchvision as tv

from net import Net
from torch.autograd import Variable


# 一些参数配置
class DefaultConfigs(object):
    lr = 0.001  # 学习率（默认值：1e-3
    epochs = 10  # 训练次，越多就越好
    num_classes = 10  # 分类
    image_size = 32  # 图片大小 ，可以改，因为用的是 resnet18 的网络，越大就越慢
    batch_size = 40  # 批量大小，看自己电脑的配置,需要占用 CPU或者GPU资源
    channels = 3  # 通道数
    use_gpu = t.cuda.is_available()  # 启用gpu，如果电脑不支持，直接设置为 False ，GPU 训练效果最好


config = DefaultConfigs()
config.use_gpu = False  # 我的电脑不支持，设置为 False

# 对Tensor进行变换 颜色转换   mean=给定均值：(R,G,B) std=方差：（R，G，B）
normalize = tv.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])


# 训练数据
def loadtraindata():
    path = "./imageData/train"  # 路径
    trainset = tv.datasets.ImageFolder(path,
                                       transform=tv.transforms.Compose([
                                           tv.transforms.Resize((config.image_size, config.image_size)),
                                           tv.transforms.CenterCrop(config.image_size),
                                           tv.transforms.ToTensor(),
                                           normalize])
                                       )

    trainloader = t.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=6)
    return trainloader


# 训练模型（支持自动GPU加速）
def train(epochs):
    trainloader = loadtraindata()
    model = Net()
    optimizer = t.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)  # 学习率为0.001
    criterion = t.nn.CrossEntropyLoss()  # 损失函数也可以自己定义，我们这里用的交叉熵损失函数
    # GPU
    if config.use_gpu:
        model = model.cuda()
        criterion = criterion.cuda()

    # 训练部分
    for epoch in range(epochs):  # 训练的数据量为5个epoch，每个epoch为一个循环
        # 每个epoch要训练所有的图片，每训练完成200张便打印一下训练的效果（loss值）
        running_loss = 0.0  # 定义一个变量方便我们对loss进行输出
        print("Epoch {}/{}".format(epoch + 1, epochs))
        for i, data in enumerate(trainloader, 0):  # 这里我们遇到了第一步中出现的trailoader，代码传入数据
            # enumerate是python的内置函数，既获得索引也获得数据
            # get the inputs
            inputs, labels = data  # data是从enumerate返回的data，包含数据和标签信息，分别赋值给inputs和labels

            # wrap them in Variable
            inputs, labels = Variable(inputs), Variable(labels)  # 转换数据格式用Variable

            optimizer.zero_grad()  # 梯度置零，因为反向传播过程中梯度会累加上一次循环的梯度

            # forward + backward + optimize
            outputs = model(inputs)  # 把数据输进CNN网络net
            loss = criterion(outputs, labels)  # 计算损失值
            loss.backward()  # loss反向传播
            optimizer.step()  # 反向传播后参数更新
            running_loss += loss.data.item()  # loss累加

            if running_loss % 20 ==0:
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 20))  # 然后再除以200，就得到这两百次的平均损失值
                running_loss = 0.0  # 这一个200次结束后，就把running_loss归零，下一个200次继续使用

    t.save(model, 'model/' + str(epochs) + "_net.pkl")  # 保存整个神经网络的结构和模型参数
    t.save(model.state_dict(), 'model/' + str(epochs) + "_net_params.pkl")  # 只保存神经网络的模型参数
    print('训练结束')


# 开始训练
if __name__ == "__main__":
    train(config.epochs)

调用代码：

import torch as t
import torchvision as tv
from PIL import Image
import matplotlib.pyplot as plt
from torch.autograd import Variable
import numpy as np
from net import Net

bCuda = t.cuda.is_available()  # 是否开启 GPU
bCuda = False  # 不启用GPU 我的电脑不支持
device = t.device("cuda:0" if bCuda else "cpu")

img_size = 32  # 图片大小，可以改

# 对Tensor进行变换 颜色转换   mean=给定均值：(R,G,B) std=方差：（R，G，B）
normalize = tv.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
transform = tv.transforms.Compose(
    [tv.transforms.Resize([img_size, img_size]), tv.transforms.CenterCrop([img_size, img_size]),
     tv.transforms.ToTensor(), normalize])

# 分类数组
classes = ['凹下标志-0', '凸上标志-1', '打滑标志-2', '左弯标志-3', '右弯标志-4', '连续转弯标志-5', '00020-6', '00021-7', '00022-8', '00023-9']


# 显示图片方法
def imshow(img):
    plt.imshow(img)
    plt.show()


# 单张图片调用
def prediect(model, img_path, imgType, isShowSoftmax=False, isShowImg=False):
    t.no_grad()
    image_PIL = Image.open(img_path)
    # imshow(image_PIL)

    image_tensor = transform(image_PIL)
    # 以下语句等效于 img = torch.unsqueeze(image_tensor, 0)
    image_tensor.unsqueeze_(0)
    # 没有这句话会报错
    image_tensor = image_tensor.to(device)
    out = model(image_tensor)
    # 得到预测结果，并且从大到小排序
    _, indices = t.sort(out, descending=True)

    # 返回每个预测值的百分数
    percentage = t.nn.functional.softmax(out, dim=1)[0] * 100

    # 是否显示每个分类的预测值
    item = indices[0]
    if isShowSoftmax:
        for idx in item:
            ss = percentage[idx]
            value = ss.item();
            name = classes[idx]
            print('名称：', name, '预测值：', value)

    # 预测最大值
    _, predicted = t.max(out.data, 1)
    maxPredicted = classes[predicted.item()]
    maxAccuracy = percentage[item[0]].item()
    if imgType == maxPredicted:
        print('预测正确，预测结果：', maxPredicted, '预测值：', maxAccuracy)
    else:
        print('预测错误，正确结果：', imgType, '，预测结果：', maxPredicted, '预测值：', maxAccuracy, '图片：', img_path)
    if isShowImg:
        plt.imshow(image_PIL)
        plt.show()


# 测试集
def loadtestdata():
    path = "./imageData/test/"
    testset = tv.datasets.ImageFolder(path, transform=transform)
    testloader = t.utils.data.DataLoader(testset, batch_size=40, shuffle=True, num_workers=6)
    return testloader


# 测试全部
def testAll(model):
    testloader = loadtestdata()

    dataiter = iter(testloader)
    images, labels = dataiter.next()
    print(labels)

    print('真实值: '
          , " ".join('%5s' % classes[labels[j]] for j in range(25)))  # 打印前25个GT（test集里图片的标签）
    outputs = model(Variable(images))
    _, predicted = t.max(outputs.data, 1)

    print('预测值: ', " ".join('%5s' % classes[predicted[j]] for j in range(25)))
    # 打印前25个预测值

    imshow2(tv.utils.make_grid(images, nrow=5))  # nrow是每行显示的图片数量，缺省值为8


def imshow2(img):
    img = img / 2 + 0.5  # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


if __name__ == '__main__':
    # 直接加载
    model = t.load('model/10_net.pkl')

    # 加载2 ，看官方的解释
    # model = Net()  # 10 分类数量
    # load_weights = t.load('model/10_net_params.pkl', map_location='cpu')
    # model.load_state_dict(load_weights)

    model = model.to(device)  # GPU
    model.eval()  # 运行模式

    # 测试全部图片
    #testAll(model)

    # 测试一张图片
    # # 凹下标志-0
    # prediect(model,'imageData/test/00000/01160_00000.png', classes[0], False, False)
    # prediect(model,'imageData/test/00000/01160_00001.png', classes[0], False, False)
    # prediect(model,'imageData/test/00000/01160_00002.png', classes[0], False, False)
    # prediect(model,'imageData/test/00000/01798_00000.png', classes[0], False, False)
    # prediect(model,'imageData/test/00000/01798_00001.png', classes[0], False, False)
    # prediect(model,'imageData/test/00000/01798_00002.png', classes[0], False, False)
    #
    # # 凸上标志-1
    # prediect(model,'imageData/test/00001/00029_00000.png', classes[1], False, False)
    # prediect(model,'imageData/test/00001/00029_00001.png', classes[1], False, False)
    # prediect(model,'imageData/test/00001/00029_00002.png', classes[1], False, False)
    # prediect(model,'imageData/test/00001/00079_00000.png', classes[1], False, False)
    # prediect(model,'imageData/test/00001/00079_00002.png', classes[1], False, False)
    # prediect(model,'imageData/test/00001/00079_00001.png', classes[1], False, False)
    #
    # # 打滑标志-2
    # prediect(model,'imageData/test/00002/01503_00000.png', classes[2], False, False)
    # prediect(model,'imageData/test/00002/01503_00001.png', classes[2], False, False)
    # prediect(model,'imageData/test/00002/01503_00002.png', classes[2], False, False)
    # prediect(model,'imageData/test/00002/01515_00000.png', classes[2], False, False)
    # prediect(model,'imageData/test/00002/01515_00001.png', classes[2], False, False)
    # prediect(model,'imageData/test/00002/01515_00002.png', classes[2], False, False)
    #
    # # 左弯标志-3
    # prediect(model,'imageData/test/00003/00207_00000.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/00207_00001.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/00207_00002.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/00211_00000.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/00211_00001.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/00211_00002.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/02664_00000.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/02664_00001.png', classes[3], False, False)
    # prediect(model,'imageData/test/00003/02664_00002.png', classes[3], False, False)
    #
    # # 右弯标志-4
    # prediect(model,'imageData/test/00004/00214_00000.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/00214_00001.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/00214_00002.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/00282_00000.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/00282_00001.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/00282_00002.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/02567_00000.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/02567_00001.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/02567_00002.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/02660_00000.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/02660_00001.png', classes[4], False, False)
    # prediect(model,'imageData/test/00004/02660_00002.png', classes[4], False, False)
    #
    # # 连续转弯标志-5
    # prediect(model,'imageData/test/00005/00575_00000.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/00575_00001.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/00575_00002.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/01893_00000.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/01893_00001.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/01893_00002.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/02225_00000.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/02225_00001.png', classes[5], False, False)
    # prediect(model,'imageData/test/00005/02225_00002.png', classes[5], False, False)
    #
    #
    # # 00020-6
    # prediect(model,'imageData/test/00020/00230_00000.png', classes[6], False, False)
    # prediect(model, 'imageData/test/00020/00230_00001.png', classes[6], True, True)
    # prediect(model,'imageData/test/00020/00230_00002.png', classes[6], False, False)
    # prediect(model,'imageData/test/00020/00231_00000.png', classes[6], False, False)
    # prediect(model,'imageData/test/00020/00231_00001.png', classes[6], False, False)
    # prediect(model,'imageData/test/00020/00231_00002.png', classes[6], False, False)
    #
    # # 00021-7
    # prediect(model, 'imageData/test/00021/00375_00000.png', classes[7], False, False)
    # prediect(model, 'imageData/test/00021/00375_00001.png', classes[7], False, False)
    # prediect(model, 'imageData/test/00021/00375_00002.png', classes[7], False, False)
    # prediect(model, 'imageData/test/00021/00478_00000.png', classes[7], False, False)
    # prediect(model, 'imageData/test/00021/00478_00001.png', classes[7], False, False)
    # prediect(model, 'imageData/test/00021/00478_00002.png', classes[7], False, False)
    #
    # # 00022-8
    # prediect(model, 'imageData/test/00022/00020_00000.png', classes[8], False, False)
    # prediect(model, 'imageData/test/00022/00020_00001.png', classes[8], False, False)
    # prediect(model, 'imageData/test/00022/00020_00002.png', classes[8], False, False)
    # prediect(model, 'imageData/test/00022/00048_00000.png', classes[8], False, False)
    # prediect(model, 'imageData/test/00022/00048_00001.png', classes[8], False, False)
    # prediect(model, 'imageData/test/00022/00048_00002.png', classes[8], False, False)
    #
    # # 00023-9
    # prediect(model, 'imageData/test/00023/00465_00000.png', classes[9], False, False)
    # prediect(model, 'imageData/test/00023/00465_00001.png', classes[9], False, False)
    # prediect(model, 'imageData/test/00023/00465_00002.png', classes[9], False, False)
    # prediect(model, 'imageData/test/00023/00535_00000.png', classes[9], False, False)
    # prediect(model, 'imageData/test/00023/00535_00001.png', classes[9], False, False)
    # prediect(model, 'imageData/test/00023/00535_00002.png', classes[9], False, False)