[PyTorch 卷积]实战自定义的图片归类

流程

1. 根据自己的定义，收集图片并归类
2. 读取图片数据和归类标签，保存数据集
3. 固定图片大小 (会变形)，归一化转张量
4. 定义超参数，损失函数和优化器等
5. 炼丹，重复查看损失值准确率等指标
6. 保存模型参数，加载测试图片分类效果

• Python 3.8
• Torch 1.9.0
• Pillow 10.0
• Torchvision
• Numpy
• Pandas
• Matplotlib

图片数据生成

``````# -*- coding: utf-8 -*-
import os
import pickle as pkl
import pandas as pd
from PIL import Image

all_cate = []
data_set = []
directory = "./data/train"
for index, data in enumerate(os.walk(directory)):
root, dirs, files = data

if index == 0:
all_cate += dirs
else:
sorted(all_cate)

root_names = root.split("\\")
dir_name = root_names[-1]

for img in files:
img_path = root + "\\" + img
img_np = Image.open(img_path)
dict = {}
dict['img_np'] = img_np
dict['label'] = all_cate.index(dir_name) + 1
data_set.append(dict)

# 字典转DataFrame
df = pd.DataFrame(data_set)
pkl.dump(df, open('data/train_dataset.p', 'wb'))
open("data/all_cate.txt", encoding="utf-8", mode="w+").write("\n".join(all_cate))

print("存档数据成功~")

这里是读取序列化的图片信息，对所有图片统一像素 (一般配置电脑最好在 100px 以内，不然会很卡) 并标准归一化后，转换为 Tensor。然后判断图片通道数，如果是灰色图，可以复制张量三次以创建三个通道，最后通过 torch 的 DataLoader 在训练前完成数据集的加载。

# -*- coding: utf-8 -*-
import torch
from torchvision import transforms
import pickle as pkl
from torch.utils.data import Dataset

class DataSet(Dataset):

def __init__(self, pkl_file):
self.dataFrame = df

def __len__(self):
return len(self.dataFrame)

def __getitem__(self, item):

img_np = self.dataFrame.iloc[item, 0]
label = self.dataFrame.iloc[item, 1]

transform = transforms.Compose([
transforms.Resize((100, 100)),  # 根据需要调整图像大小
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])    # 标准归一化, p1.均值  p2.方差
])
image_tensor = transform(img_np)

if image_tensor.shape[0] == 1:
image_tensor = image_tensor.repeat(3, 1, 1)

res = {
'img_tensor': image_tensor,
'label': torch.LongTensor([label-1])    # 需要实际的索引值
}

return res``````

神经网络模型

``````# -*- coding: utf-8 -*-
import torch.nn as nn
import torch
import math
import torch.functional as F

class CNN(nn.Module):

def __init__(self):
super(CNN, self).__init__()

self.layer1 = nn.Sequential(
nn.Conv2d(3, 25, kernel_size=3),
nn.BatchNorm2d(25),
nn.ReLU(inplace=True)
)

self.layer2 = nn.Sequential(
nn.MaxPool2d(kernel_size=2, stride=2)
)

self.layer3 = nn.Sequential(
nn.Conv2d(25, 50, kernel_size=3),
nn.BatchNorm2d(50),
nn.ReLU(inplace=True)
)

self.layer4 = nn.Sequential(
nn.MaxPool2d(kernel_size=2, stride=2)
)

self.fc = nn.Sequential(
nn.Linear(50 * 23 * 23, 1024),
nn.ReLU(inplace=True),
nn.Linear(1024, 128),
nn.ReLU(inplace=True),
nn.Linear(128, 5)
)

def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = x.view(x.size(0), -1)

x = self.fc(x)

return x``````

开始训练

``````# -*- coding:utf-8 -*-
import torch
import matplotlib.pyplot as plt
from data_set import DataSet
from utils import *
import cnn
import torch.nn as nn
import numpy as np
import torch.optim as optim

# 定义超参数
batch_size = 1
learning_rate = 0.02
num_epoches = 1

# 加载图片tensor训练集
tain_dataset = DataSet("data/train_dataset.p")

model = cnn.CNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 训练模型
train_loses = []
records = []
for i in range(num_epoches):
img = data['img_tensor']
label = data['label'].view(-1)

out = model(img)
loss = criterion(out, label)
train_loses.append(loss.data.item())
loss.backward()
optimizer.step()

if ii % 50 == 0:
print('epoch: {}, loop: {}, loss: {:.4}'.format(i, ii, np.mean(train_loses)))

records.append([np.mean(train_loses)])

# 绘制模型的损失，准确率走势图
train_loss = [data[0] for data in records]
plt.plot(train_loss, label = 'Train Loss')
plt.xlabel('Steps')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 模型评估(略)
# model.eval()

# 模型保存
torch.save(model, 'params/cnn_imgs_02.pkl')``````

模型检测

``````# -*- coding:utf-8 -*-
import torch
import matplotlib.pyplot as plt
from data_set import DataSet
from utils import *
import torchvision
from PIL import Image
from torchvision import transforms
import cnn

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

img_path= "imgs/05.jpg"
img_np = Image.open(img_path)
transform = transforms.Compose([
transforms.Resize((100, 100)),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
image_tensor = transform(img_np)

# 如果是灰度图片
if image_tensor.shape[0] == 1:
image_tensor = image_tensor.repeat(3, 1, 1)

image_tensor = image_tensor.view(-1, 3, 100, 100)

predict = model(image_tensor)
indices = torch.max(predict, 1)[1].item()

all_cate = []
for line in open("data/all_cate.txt", encoding="utf-8", mode="r"):
all_cate.append(line.strip())

cate_name = ""
try:
cate_name = all_cate[indices]
except ValueError:
cate_name = "未知"

print("识别结果是：", cate_name)
# imshow(torchvision.utils.make_grid(image_tensor))
# 原图显示
img_np.show()
exit()``````

(=￣ω￣=)··· 暂无内容！

UP @ 公众号： ZERO开发

50

1

6

4