PyTorch 基本图像分类器

• 入门深度学习，手里刚刚有点算力，就想训练一下自己的数据。

• 从最基本的工作开始做起：

• 本文采用了最朴素的卷积+全链接的形式训练一个三类分类器，并进行可视化

• 直接分享出代码

"""
# 分类任务
# 导入自己的数据 进行训练和测试

0 胶带
1 芦荟胶
2 魔方
"""

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import os
import collections  # 序列化列表

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

"""
整体结构：
1.加载数据
2.定义网络结构
3.训练网络参数
4.用测试集测试
"""

# 对类别进行对应编号
class_list = [
    "tape",
    'makeup',
    'cude',
    'background']


# 有序化
class_list = {x:class_list[x] for  x in range(4)}

print('class_list: ', class_list)

print(type(class_list))

""" 
数据变换
"""
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor()])

"""
加载数据
"""
batch_size_train = 20
batch_size_test = 20

data_root = os.path.abspath("../../../00.Dataset/my_data")  # get data root path

print(data_root)

# 对于单纯的分类任务，使用ImageFolder简化数据加载
mnist_data_train = torchvision.datasets.ImageFolder(data_root, transform=transform)

train_loader = torch.utils.data.DataLoader(mnist_data_train,
                                           batch_size_train, shuffle=True, num_workers=4)

mnist_data_test = torchvision.datasets.ImageFolder(data_root, transform=transform)

test_loader = torch.utils.data.DataLoader(mnist_data_test,
                                          batch_size_test, shuffle=True, num_workers=4)



# 搭建网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 重定义函数
        self.conv1 = nn.Conv2d(3, 7, 5, 1)  # kernel_1 in-1  out-6  kernel_size=5x5  ->24x24
        self.bn1 = nn.BatchNorm2d(7)  # 输入通道数
        self.conv2 = nn.Conv2d(7, 20, 5, 1)
        self.pool = nn.MaxPool2d(2, 2)  # max pooling
        self.liner1 = nn.Linear(4 * 4 * 20, 20)  # 需要计算才能知道
        self.liner2 = nn.Linear(20, 4)  # in_features, out_features, bias=True

    # 自动被调用
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # Conv -> Relu -> Pool  -> 12x12 @ 6 CHs
        self.bn1(x)
        x = self.pool(F.relu(self.conv2(x)))  # -> 8x8 @6 ->  4x4 @ 20 CHs
        # print("x_sizer: ", x.shape)

        x = x.view(-1, 4 * 4 * 20)
        x = F.relu(self.liner1(x))
        x = self.liner2(x)

        # return x.log_softmax(x, dim=1)
        return x


net = Net().cuda()

"""
损失函数
优化器
"""

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

"""
训练集训练
"""


def train():
    for epoch in range(500):  # 遍历两次
        running_loss = 0
        for i, data in enumerate(train_loader):  # 遍历一遍全部的数据集
            inputs, labels = data  # batch_size个样本
            optimizer.zero_grad()

            # 前向传播
            output = net(inputs.cuda())

            # print("out : ", output)
            # print("labels : ", labels)

            loss = criterion(output, labels.cuda())
            # 反向传播
            loss.backward()
            optimizer.step()  # 更新变量

            running_loss += loss.item()
            if (i % 20 == 0):
                print("epoch:{} , i:{} , running_loss:{}".format(epoch, i, running_loss))
                running_loss = 0

    print("Finished Training!")


# train()
#
# # 仅保存模型参数  注意 训练完保存一次，再次运行务必注释掉该句
# print("Saving model paremeters...")
# torch.save(net.state_dict(), "../model/params_mydata.pkl")

# 加载模型参数
net.load_state_dict(torch.load("../model/params_mydata.pkl"))

# 保存整个模型
# torch.save(net, "../data/MNIST/model.pkl")
# 加载模型
# net = torch.load("../data/MNIST/model.pkl")

# Test
"""
img
predicted
img_size   
"""


def imshow(img, nrow, predicted, img_size=28):
    # img = img / 2 + 0.5  # unnormalize  显示用
    npimg = img.numpy()
    # print("npimg shape: ", npimg.shape)
    npimg_trans = np.transpose(npimg, (1, 2, 0))  # 调换了位置
    # print("npimg_trans shape: ", npimg_trans.shape) #
    for i in range(len(predicted) // nrow):  # 行
        for j in range(nrow):  # 列
            # plt.text(5 + (img_size+3) * j, 5 + (img_size+3) * i, "pred:" + str(predicted[i * nrow + j].item()), color="r") #
            plt.text(5 + (img_size+3) * j, 5 + (img_size+3) * i, class_list[predicted[i * nrow + j].item()], color="r") #

    plt.imshow(npimg_trans)
    plt.show()


# 拿出一个batch来测试

test_data1 = iter(test_loader)
images, labels = test_data1.next()  # 取出其中的一个batch  里面有4个样本
images2, labels2 = test_data1.next()  # 取出其中的一个batch  里面有4个样本

outputs = net(images.cuda())
_, predicted = torch.max(outputs, 1)

# 预测
print('Predicted: ', ' '.join('%3s' % int(predicted[j]) for j in range(batch_size_test)))
# 实际
print("Grounf truth: ", labels)


nrow = 5
imshow(torchvision.utils.make_grid(images, nrow=nrow), nrow, predicted)

# 测试 全部样本上的整体准确度
correct = 0
total = 0

class_correct = list(0. for i in range(batch_size_test))
class_total = list(0 for j in range(batch_size_test))

with torch.no_grad():  # Dont need backward
    for data in test_loader:
        images, labels = data
        outputs = net(images.cuda())
        _, predicted = torch.max(outputs, dim=1)  # 每行的最大值  可以推出，列是对不同变量，行内是各类的概率
        total += labels.size(0)
        correct += (predicted == labels.cuda()).sum().item()

        c = (predicted == labels.cuda()).squeeze()

        for i in range(batch_size_test):  # 每个 batch 内的样本个数
            label = labels[i]
            class_total[label] += 1
            class_correct[label] += c[i].item()

print("Total samples: %4d Accuracy: %3f %%" % (total, 100 * correct / total))
print("Accuracy on each class:")
# 在我们这，标签的索引就是类别 ！ Take it easy
for i in range(10):
    print("Accuracy each class:%2d  %3f" % (i, class_correct[i] / class_total[i]))

• 结果展示
• 分别对应： 0 胶带； 1 芦荟胶 ； 2 魔方 ； 3 背景

Total samples:  400 Accuracy: 100.000000 %
Accuracy on each class:
Accuracy each class: 0  1.000000
Accuracy each class: 1  1.000000
Accuracy each class: 2  1.000000
Accuracy each class: 3  1.000000

• epoch小的时候，准确率还是很低的，但我迭代上几百次也没关系阿hehe。epoch是第一生产力！

• 分类任务就做到这，接下来再进行下一步的工作