pytorch实现猫狗分类+数据集
PyTorch实现猫狗分类代码
实现环境：python3.7，pytorch，jupyter notebook
数据集： Kaggle网站猫狗分类数据集，训练集包括12500张猫和12500张狗，测试集为12500张猫狗无标签图片，我们需要把训练集的猫狗照片各取出2500张作为验证集。
开始部分代码：导入一些必要的包，os包集成了一些对文件路径及目录进行操作的类，time包是一些和时间相关的方法
import torch
import torchvision
from torch.autograd import Variable
from torchvision import datasets, transforms
import os
import matplotlib.pyplot as plt
%matplotlib inline
import time
数据的载入：读取数据并载入
data_dir = 'DogsVSCats'
data_transform = {x:transforms.Compose([transforms.Scale([64, 64]),
transforms.ToTensor()]) for x in ['train', 'valid']}
image_datasets = {x:datasets.ImageFolder(root = os.path.join(data_dir,x),
transform = data_transform[x]) for x in ['train', 'valid']}
dataloader = {x:torch.utils.data.DataLoader(dataset = image_datasets[x],
batch_size = 16,
shuffle = True) for x in ['train', 'valid']}
数据预览：
# 数据预览
X_example, Y_example = next(iter(dataloader['train']))
print(u'X_example个数{}'.format(len(X_example)))
print(u'Y_example个数{}'.format(len(Y_example)))
index_classes = image_datasets['train'].class_to_idx
print(index_classes)
example_classes = image_datasets['train'].classes
print(example_classes)
img = torchvision.utils.make_grid(X_example)
img = img.numpy().transpose([1,2,0])
print([example_classes[i] for i in Y_example])
plt.imshow(img)
plt.show()
结果：
X_example个数16
Y_example个数16
{‘Cat’: 0, ‘Dog’: 1}
[‘Cat’, ‘Dog’]
[‘Cat’, ‘Cat’, ‘Dog’, ‘Dog’, ‘Dog’, ‘Cat’, ‘Cat’, ‘Dog’, ‘Dog’, ‘Dog’, ‘Cat’, ‘Cat’, ‘Dog’, ‘Cat’, ‘Dog’, ‘Cat’]
构建VGGnet
# 模型搭建 简化了的VGGnet
class Models(torch.nn.Module):
def __init__(self):
super(Models,self).__init__()
self.Conv = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.Conv2d(64, 64, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size = 2, stride = 2),
torch.nn.Conv2d(64, 128, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.Conv2d(128, 128, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size = 2, stride = 2),
torch.nn.Conv2d(128, 256, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.Conv2d(256, 256, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.Conv2d(256, 256, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size = 2, stride = 2),
torch.nn.Conv2d(256, 512, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(kernel_size = 2, stride = 2))
self.Classes = torch.nn.Sequential(
torch.nn.Linear(4*4*512, 1024),
torch.nn.ReLU(),
torch.nn.Dropout(p = 0.5),
torch.nn.Linear(1024, 1024),
torch.nn.ReLU(),
torch.nn.Dropout(p = 0.5),
torch.nn.Linear(1024, 2))
def forward(self, inputs):
x = self.Conv(inputs)
x = x.view(-1, 4*4*512)
x = self.Classes(x)
return x
model = Models()
print(model)
模型训练
loss_f = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr = 0.00001)
Use_gpu = torch.cuda.is_available()
if Use_gpu:
model = model.cuda()
epoch_n = 10
time_open = time.time()
for epoch in range(epoch_n):
print('epoch{}/{}'.format(epoch,epoch_n-1))
print('-'*10)
for phase in ['train','valid']:
if phase == 'train':
print('training...')
model.train(True)
else:
print('validing...')
model.train(False)
running_loss = 0.0
running_corrects = 0.0
for batch,data in enumerate(dataloader[phase],1):
X,Y = data
X,Y = Variable(X).cuda(),Variable(Y).cuda()
y_pred = model(X)
_,pred = torch.max(y_pred.data,1)
optimizer.zero_grad()
loss = loss_f(y_pred,Y)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.data.item()
running_corrects += torch.sum(pred == y.data)
if batch%500 == 0 and phase =='train':
print('batch{},trainLoss;{:.4f},trainAcc:{:.4f}'.format(batch,running_loss/bath,100*running_corrects/(16*batch)))
epoch_loss = running_loss *16/len(image_datasets[phase])
epoch_acc = 100*running_corrects/len(image_datasets[phase])
print('{} Loss:{:.4f} Acc:{:.4f}%'.format(phase,epoch_loss,epoch_acc))
time_end = time.time()-time_open
print(time_end)
补充：验证集划分程序，并将划分后的猫狗类别区分（略辣鸡）
import os
import random
import shutil
from tqdm import tqdm
def move(src_path,cat_path,dog_path):
All_files = os.listdir(src_path) #读取原路径下的所有文件，即图片
length = len(All_files)
cat,dog,i = 0,0,0
while i < 5000:
number = random.randint(0,length) #生成随机数
img_name = All_files[number] #从图片中随机选择一个图片
name = img_name.split('.')[0] #根据图片名按 '.'分成段，如cat.1.jpg，即'cat'，'1','jpg' 0对应'cat',以此类推
try: #加报警检测机制，如果有重复的，不报错直接跳过
if name == 'cat':
if cat < 2500:
shutil.move(src_path + '\\' +img_name, cat_path) #剪切到目标路径
cat = cat + 1
length = length - 1 #成功剪切一次后，总数减1
i = i + 1
else:
if dog < 2500:
shutil.move(src_path + '\\' +img_name, dog_path)
dog = dog + 1
length = length - 1
i = i + 1
if cat == 2500 & dog ==2500:
break
except shutil.Error:
i = i - 1 #若剪切不成功，则此次不加1
continue
if __name__ == "__main__":
move("D:\Data\\train", #原目标文件夹
"D:\Data\cat", #猫
"D:\Data\dog") #狗
备注：
猫狗分类数据集：
链接：https://pan.baidu.com/s/1pNKPcOb062npE85NFX_ccw
提取码：tfh0