CIFAR10数据集训练及测试2
一、数据集解析
#数据读取与解析
import pickle
import numpy as np
import cv2
def unpickle(file):
with open(file,'rb') as fo:
dict = pickle.load(fo,encoding='bytes')
return dict
label_name = ["airplane",
"automobile",
"bird",
"cat",
"deer",
"dog",
"frog",
"horse",
"ship",
"truck"]
import glob
import numpy
import os
#解析训练集
im_train_list = glob.glob("G:\pycharm-work\CIFAR10\data_batch_*")
print(im_train_list)
save_path = "G:\pycharm-work\CIFAR10\TRAIN"
for l in im_train_list:
print(l)
l_dict = unpickle(l)
print(l_dict.keys())
for im_idx , im_data in enumerate(l_dict[b'data']):
# print(im_idx)
# print(im_data)
im_label = l_dict[b'labels'][im_idx]
im_name = l_dict[b'filenames'][im_idx]
# print(im_label,im_name)
im_label_name = label_name[im_label]
#对数据进行reshape 则需要转为numpy格式
im_data = np.reshape(im_data,[3,32,32])
im_data = np.transpose(im_data,(1,2,0))
# cv2.imshow("im_data",cv2.resize(im_data,(200,200)))
# cv2.waitKey(0)
if not os.path.exists("{}/{}".format(save_path,im_label_name)):
os.mkdir("{}/{}".format(save_path,im_label_name))
cv2.imwrite("{}/{}/{}".format(save_path,im_label_name,
im_name.decode("utf-8")),im_data)
#解析训练集
im_test_list = glob.glob("G:\pycharm-work\CIFAR10\test_batch_*")
print(im_test_list)
save_path = "G:\pycharm-work\CIFAR10\TEST"
for l in im_test_list:
print(l)
l_dict = unpickle(l)
print(l_dict.keys())
for im_idx , im_data in enumerate(l_dict[b'data']):
# print(im_idx)
# print(im_data)
im_label = l_dict[b'labels'][im_idx]
im_name = l_dict[b'filenames'][im_idx]
# print(im_label,im_name)
im_label_name = label_name[im_label]
#对数据进行reshape 则需要转为numpy格式
im_data = np.reshape(im_data,[3,32,32])
im_data = np.transpose(im_data,(1,2,0))
# cv2.imshow("im_data",cv2.resize(im_data,(200,200)))
# cv2.waitKey(0)
if not os.path.exists("{}/{}".format(save_path,im_label_name)):
os.mkdir("{}/{}".format(save_path,im_label_name))
cv2.imwrite("{}/{}/{}".format(save_path,im_label_name,
im_name.decode("utf-8")),im_data)
二、自定义数据集加载
import glob
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import DataLoader,Dataset
from torch.utils.tensorboard import SummaryWriter
import time
import os
from PIL import Image
import numpy as np
import cv2
label_name = ["airplane",
"automobile",
"bird",
"cat",
"deer",
"dog",
"frog",
"horse",
"ship",
"truck"]
label_dict = {}
for idx, name in enumerate(label_name):
label_dict[name] = idx
# print(label_dict)
def default_loader(path):
return Image.open(path).convert("RGB")
train_transform = transforms.Compose([
transforms.RandomResizedCrop((28,28)),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.RandomRotation(90),
transforms.RandomGrayscale(0.1),
transforms.ColorJitter(0.3,0.3,0.3,0.3),
transforms.ToTensor()
])
class MyDataset(Dataset):
def __init__(self,im_list,transform = None,loader = default_loader):
super(MyDataset,self).__init__()
imgs = []
for im_item in im_list:
# print(im_item)
im_label_name = im_item.split("\\")[-2]
imgs.append([im_item,label_dict[im_label_name]])
self.imgs =imgs
self.transform =transform
self.loader = loader
def __getitem__(self, index):
im_path, im_label = 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_label
def __len__(self):
return len(self.imgs)
im_train_list = glob.glob("G:\pycharm-work\CIFAR10\TRAIN\*\*.png")
im_test_list = glob.glob("G:\pycharm-work\CIFAR10\TEST\*\*.png")
train_data_set = MyDataset(im_train_list,transform=train_transform)
test_data_set = MyDataset(im_test_list,transform = transforms.ToTensor())
# train_data_loader = DataLoader(dataset=train_data_set,batch_size=6,shuffle=True,num_workers=4)
# test_data_loader = DataLoader(dataset=test_data_set,batch_size=6,shuffle=False,num_workers=4)
train_data_loader =DataLoader(dataset=train_data_set,batch_size=128,shuffle=True)
test_data_loader = DataLoader(dataset=test_data_set,batch_size=128,shuffle=False)
print("num_of_train",len(train_data_set))
print("num_of_test",len(test_data_set))三、搭建神经网络模型
class VGGBase(nn.Module):
def __init__(self):
super(VGGBase, self).__init__()
# 3 * 28 * 28
self.conv1 = nn.Sequential(
nn.Conv2d(3,64,kernel_size=3,stride=1,padding=1),
nn.BatchNorm2d(64),
nn.ReLU()
)
#14*14
self.max_pooling1 = nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
self.conv2_1 = nn.Sequential(
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(128),
nn.ReLU()
)
self.conv2_2 = nn.Sequential(
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(128),
nn.ReLU()
)
# 7*7
self.max_pooling2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.conv3 = nn.Sequential(
nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(256),
nn.ReLU()
)
self.conv3_2 = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(256),
nn.ReLU()
)
self.max_pooling3 = nn.MaxPool2d(kernel_size=2, stride=2, padding=1)# 获得4 * 4
self.conv4 = nn.Sequential(
nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(512),
nn.ReLU()
)
self.conv4_2 = nn.Sequential(
nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(512),
nn.ReLU()
)
self.max_pooling4 = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc = nn.Linear(512 * 4,10)
def forward(self,x):
batchsize = x.size(0)
out = self.conv1(x)
out = self.max_pooling1(out)
out = self.conv2_1(out)
out = self.conv2_2(out)
out = self.max_pooling2(out)
out = self.conv3(out)
out = self.conv3_2(out)
out = self.max_pooling3(out)
out = self.conv4(out)
out = self.conv4_2(out)
out = self.max_pooling4(out)
out = out.view(batchsize, -1)
out = self.fc(out)
out = F.log_softmax(out,dim=1)
return out
def VGGNet():
return VGGBase()四、训练与测试
# 定义训练设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net =VGGNet()
net = net.to(device)
# 损失函数
loss_fn =nn.CrossEntropyLoss()
loss_fn =loss_fn.to(device)
# 优化器
learning_rate =0.01
#weight_decay 正则项 momentum 动量
# optimizer =torch.optim.SGD(net.parameters(),lr = learning_rate,momentum=0.9,weight_decay=5e-4)
optimizer =torch.optim.Adam(net.parameters(),lr = learning_rate)
#学习率调整 指数衰减
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.9)
# 添加Tensorboard
if not os.path.exists("log"):
os.mkdir("log")
writer =SummaryWriter("log")
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 记录测试的次数
total_test_step = 0
epoch = 30
batch_size = 128
start_time = time.time()
step_n = 0
for epochidx in range(epoch):
print("-----第{}轮训练开始------".format(epochidx + 1))
# 训练步骤开始
net.train() #BN与dropout 更新
for i,data in enumerate(train_data_loader):
imgs,target =data
imgs = imgs.to(device)
target = target.to(device)
output = net(imgs)
loss = loss_fn(output, target)
#优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
# total_train_step +=1
# if total_train_step % 100 == 0:
# end_time =time.time()
# print(end_time-start_time)
# print("训练次数{}, Loss:{}".format(total_train_step, loss.item()))
# writer.add_scalar("train_loss",loss.item(),total_train_step)
_,pred = torch.max(output.data,dim= 1)
correct = pred.eq(target.data).cpu().sum()
print("train epoch is ",epochidx)
print("trainlr is ",optimizer.state_dict()["param_groups"][0]["lr"])
print("train step",i,"loss is:",loss.item(),
"mini_batch correct is:",100.0 * correct / batch_size)
#记录
writer.add_scalar("train loss",loss.item(),global_step=step_n)
writer.add_scalar("train correct", 100.0 * correct / batch_size,global_step=step_n)
im = torchvision.utils.make_grid(imgs)
writer.add_image("train img",im,global_step=step_n)
step_n +=1
if not os.path.exists("models"):
os.mkdir("models")
torch.save(net.state_dict(),"models/{}.pth".format(epochidx+1))
#每个epoch后更新学习率
scheduler.step()
#对模型进行测试
sum_loss = 0
sum_correct = 0
for j, data in enumerate(test_data_loader):
net.eval() # 测试
imgs, target = data
imgs = imgs.to(device)
target = target.to(device)
output = net(imgs)
loss = loss_fn(output, target)
_, pred = torch.max(output.data, dim=1)
correct = pred.eq(target.data).cpu().sum()
sum_loss += loss.item()
sum_correct +=correct.item()
im = torchvision.utils.make_grid(imgs)
writer.add_image("test img", im, global_step=step_n)
test_loss =sum_loss*1.0/len(test_data_loader)
test_correct = sum_correct * 100.0 / len(test_data_loader)/batch_size
# 记录
writer.add_scalar("test loss", test_loss, global_step=j+1)
writer.add_scalar("test correct", test_correct, global_step=j+1)
print("epoch is", j+1, "loss is:", test_loss,
"test correct is:", test_correct)
writer.close()