基于Pytorch的cifar-10图像分类问题代码实现
之前在学习深度学习图片分类任务的时候，跟着老师的讲解实现了一个Cifair-10的图像分类任务。
数据集地址：网盘地址需要解码。
数据有50000张训练图片和10000张测试图片。
下载好数据后，在文件夹下新建两个文件夹一个为Train，一个为Test，用来保存解码后的图片。如图：
解码后的训练集：
文件的解码方式官方已经给出：
def unpickle(file):
with open(file, 'rb') as fo:
dict = pickle.load(fo, encoding="bytes")
return dict
解码训练集的py文件的具体代码如下：
import os
import pickle
import glob #文件匹配
import cv2
import numpy as np
def unpickle(file): #图片的解码方式，数据集已经给出
with open(file, 'rb') as fo:
dict = pickle.load(fo, encoding="bytes")
return dict #返回一个字典
lable_name = ["airplane",
"automobile",
"bird",
"cat",
"deer",
"dog",
"frog",
"horse",
"ship",
"truck"]
train_list = glob.glob("D:\*\cifar-10-python\cifar-10-batches-py\data_batch_*") #文件的位置
print(train_list)
save_path = "D:\*\cifar-10-python\cifar-10-batches-py\Train"#解码成图片后文件的保存位置
for l in train_list:#读去每个文件，解码
print(l)
l_dict = unpickle(l)
print(l_dict)
print(l_dict.keys())
for im_idx, im_data in enumerate(l_dict[b'data']):#enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标，一般用在 for 循环当中
print(im_idx)
print(im_data)
解码测试集的py文件的具体代码如下：
import os
import pickle
import cv2
import numpy as np
def unpickle(file): #图片的解码方式，数据集已经给出
with open(file, 'rb') as fo:
dict = pickle.load(fo, encoding="bytes")
return dict #返回一个字典
lable_name = ["airplane",
"automobile",
"bird",
"cat",
"deer",
"dog",
"frog",
"horse",
"ship",
"truck"]
import glob #文件匹配
train_list = glob.glob("D:\genglijia\cifar-10-python\cifar-10-batches-py\Test_batch")
print(train_list)
save_path = "D:\genglijia\cifar-10-python\cifar-10-batches-py\Test"
for l in train_list:#读去每个文件，解码
print(l)
l_dict = unpickle(l)
print(l_dict)
print(l_dict.keys())
for im_idx, im_data in enumerate(l_dict[b'data']):#enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标，一般用在 for 循环当中
print(im_idx)
print(im_data)
im_lable = l_dict[b'labels'][im_idx]
im_name = l_dict[b'filenames'][im_idx]
print(im_lable, im_name, im_data)
im_lable_name = lable_name[im_lable]
im_data = np.reshape(im_data, [3, 32, 32])#这里的图片的大小转换不是很清楚
im_data = np.transpose(im_data, (1, 2, 0))#更改图片维度的位置，类似矩阵的转置
if not os.path.exists("{}/{}".format(save_path,#判断path对应文件或目录是否存在，返回布尔类型
im_lable_name)):
os.mkdir("{}/{}".format(save_path,
im_lable_name))
cv2.imwrite("{}/{}/{}".format(save_path,
im_lable_name,
im_name.decode("utf-8")),
im_data)
加载本地数据集的py文件：
import glob
from torchvision import transforms #transforms用于后面数据增强
from torch.utils.data import DataLoader,Dataset#数据加载和数据读取相关的类
import os
from PIL import Image #用于图片数据的处理，有点类似opencv,数据会以numpy的形式进行存储
import numpy as np
lable_name = ["airplane","automobile","bird",
"cat","deer","dog","frog",
"horse","ship","truck"]
lable_dict = {}
#将字符串lable全部转换为数字
for idx, name in enumerate(lable_name):
lable_dict[name] = idx
print(lable_dict)
def default_loader(path):
return Image.open(path).convert("RGB")
#数据增强
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
test_transform = transforms.Compose([
transforms.CenterCrop((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
class MyDataset(Dataset):
#初始化函数
def __init__(self, im_list, transform=None, loader=default_loader): #im_list拿到当前文件夹下所有文件的一个列表,也可以将文件的路径放到一个text中
super(MyDataset, self).__init__()
imgs = []
for im_item in im_list: #拿到的其实是“\genglijia\cifar-10-python\cifar-10-batches-py\Test\airplane\*.png”
im_lable_name = im_item.split("\\")[-2]
imgs.append([im_item, lable_dict[im_lable_name]])
self.imgs = imgs#图片中的每一个元素
self.transform = transform#两个方法 transform数据增强
self.loader = loader
#读取图片数据中元素的方法
def __getitem__(self, index):
im_path,im_lable = self.imgs[index]
im_data = self.loader(im_path)
if self.transform is not None:
im_data = self.transform(im_data)
return im_data, im_lable
#计算样本的数量
def __len__(self):
return len(self.imgs)
im_train_list = glob.glob("D:\*\cifar-10-python\cifar-10-batches-py\Train\*\*.png")
im_test_list = glob.glob("D:\*\cifar-10-python\cifar-10-batches-py\Test\*\*.png")
train_dataset = MyDataset(im_train_list, transform=train_transform)
test_dataset = MyDataset(im_test_list, transform=transforms.ToTensor())
train_loader = DataLoader(dataset=train_dataset,
batch_size=64,
shuffle=True,
num_workers=0)
test_loader = DataLoader(dataset=test_dataset,
batch_size=64,
shuffle=False,
num_workers=0)#num_workers类似于多个进程同时对数据加载
print("num_of_train", len(train_dataset))
print("num_of_test", len(test_dataset))
定义网络结构的py文件（用的是经典resnet残差网络机构，也可以用其他的网络结构例如：vggnet、mobilenet等）：
import torch
import torch.nn as nn
import torch.nn.functional as F
class ResBlock(nn.Module):
def __init__(self, in_channel, out_channel, stride=1):
super(ResBlock, self).__init__()
self.layer = nn.Sequential(
nn.Conv2d(in_channel, out_channel,
kernel_size=3, stride=stride, padding=1),
nn.BatchNorm2d(out_channel),
nn.ReLU(),
nn.Conv2d(out_channel, out_channel,
kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(out_channel),
)
self.shortcut = nn.Sequential()
if in_channel != out_channel or stride > 1:
self.shortcut = nn.Sequential(
nn.Conv2d(in_channel, out_channel,
kernel_size=3, stride=stride, padding=1),
nn.BatchNorm2d(out_channel),
)
def forward(self, x):
out1 = self.layer(x)
out2 = self.shortcut(x)
out = out1 + out2
out = F.relu(out)
return out
class ResNet(nn.Module):
def make_layer(self, block, out_channel, stride, num_block):
layers_list = []
for i in range(num_block):
if i == 0:
in_stride = stride
else:
in_stride = 1
layers_list.append(block(self.in_channel,out_channel, in_stride))
self.in_channel = out_channel
return nn.Sequential(*layers_list)
def __init__(self, ResBlock):
super(ResNet, self).__init__()
self.in_channel = 32
self.conv1 = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(32),
nn.ReLU()
)
self.layer1 = \
self.make_layer(ResBlock, 64, 2, 2)
self.layer2 = \
self.make_layer(ResBlock, 128, 2, 2)
self.layer3 = \
self.make_layer(ResBlock, 256, 2, 2)
self.layer4 = \
self.make_layer(ResBlock, 512, 2, 2)
self.fc = nn.Linear(512, 10)
def forward(self,x):
out = self.conv1(x)
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 2)
out = out.view(out.size(0), -1)
out = self.fc(out)
return out
def resnet():
return ResNet(ResBlock)
训练和测试的py文件：
import torch
import torch.nn as nn
import torchvision
from resnet import resnet
from load_cifar10 import train_loader,test_loader
import os
#判断是否存在gpu
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
epoch_num = 1
lr = 0.01
batch_size = 128
net = resnet()#如果有GPU的话在后面加.to(device)
#loss
loss_func = nn.CrossEntropyLoss()
#optimizer
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
#optimizer = torch.optim.SGD(net.parameters(), lr=lr,
# momentum=0.9,weight_decay=5e-4)
#更改学习率，采用指数的方式，也可以用固定的学习率
#step_size:每5次epoch更改学习率 gamma：更改为上次学习率的0.9倍
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)
if __name__ == '__main__': #解决windows下的报错，因为使用了num_work,具体为啥我也不太清楚
for epoch in range(epoch_num):
net.train()#表明当前网络为训练的过程 train BN dropout
#如果在网络层定义了Batchnorm层则需要用net.train
#如果在网络层定义了dropout层则需要用net.eval()
for i, data in enumerate(train_loader):
inputs, labels = data
# inputs, labels = inputs.to(device), labels.to(device) #转到GPU训练
outputs = net(inputs)
loss = loss_func(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_size = inputs.size(0)
_, pred = torch.max(outputs.data, dim=1)
correct = pred.eq(labels.data).sum()
print("train step", i, "loss is:", loss.item(), "mini-batch correct is:", 1.0 * correct / batch_size)
#用来保存训练后的模型
if not os.path.exists("models"):
os.mkdir("models")
torch.save(net.state_dict(),"models/{}.pth".format(epoch+1))
scheduler.step()#更新学习率
sum_loss = 0
sum_correct = 0
for i, data in enumerate(test_loader):
net.eval()
inputs, labels = data
# inputs, labels = inputs.to(device), labels.to(device)
outputs = net(inputs)
loss = loss_func(outputs, labels)#测试集不再进行反向传播
_, pred = torch.max(outputs.data, dim=1)
correct = pred.eq(labels.data).sum()
sum_loss += loss.item()
sum_correct += correct.item()
im = torchvision.utils.make_grid(inputs)
test_loss = sum_loss*1.0/len(test_loader)
test_correct = sum_correct *1.0/len(test_loader)/batch_size
print("epoch", epoch+1, "loss is:", test_loss, "mini-batch correct is:", test_correct)
测试的结果：
这里面epoch为1，因为没有GPU训练的太慢就只训练了一次，正确率也有了百分之70左右，有条件的话多训练几次应该会达到更高。也可以改网络结构啥的，方法很多。
完成后可以在Test的文件夹中看到已经分类好的图片。
**有什么疑问可以在下面评论哦~**大家一起加油学习！！！