Pytorch框架下的Kaggle猫狗识别

Kaggle猫狗识别，训练集为25000张打好标签的猫狗图片，测试集为划分出来的5000张猫和狗的照片。

使用的网络为Le-Net，其结构图如下

 

 

代码如下:

import torch
import numpy
import matplotlib.pylab as plt
from torch.autograd import Variable
from torchvision import transforms,datasets
import cv2
import torch.nn as nn
from tensorboardX import SummaryWriter


print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.get_device_capability(),torch.cuda.get_device_name())


image_transform = transforms.Compose([
    transforms.Resize(84),
    transforms.CenterCrop(84),
    transforms.ToTensor(),
    transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
])

train_dataset = datasets.ImageFolder(root =r'D:\Python\project\venv\Include\kaggle数据集\train/',transform=image_transform)
train_loader = torch.utils.data.DataLoader(train_dataset,4,True)

test_dataset = datasets.ImageFolder(root = r'D:\Python\project\venv\Include\kaggle数据集\selftest/',transform=image_transform)
test_loader = torch.utils.data.DataLoader(test_dataset,4,True)

use_dataset = datasets.ImageFolder(root = r'D:\Python\project\venv\Include\kaggle数据集\use/',transform=image_transform)
use_loader = torch.utils.data.DataLoader(use_dataset,1,True)


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = nn.Sequential(nn.Conv2d(3,6,5),nn.ReLU(),nn.MaxPool2d(2,2))
        self.conv2 = nn.Sequential(nn.Conv2d(6,16,5),nn.ReLU(),nn.MaxPool2d(2,2))
        self.fc1 = nn.Sequential(nn.Linear(16*18*18,1024),nn.ReLU())
        self.fc2 = nn.Sequential(nn.Linear(1024,512),nn.ReLU(),nn.Linear(512,2))

    def forward(self,input):
        x = self.conv1(input)
        x = self.conv2(x)
        x = x.view(-1,16*18*18)
        x = self.fc1(x)
        x = self.fc2(x)
        x = torch.sigmoid(x)

        return x

net = Net()
net.cuda()
net.load_state_dict(torch.load(r'D:\Python\project\venv\Include\Nural_net/torch_net10(74%)'))
writer = SummaryWriter('figure')
opti = torch.optim.SGD(net.parameters(),lr = 0.0005,momentum=0.9)
loss_fun = nn.CrossEntropyLoss()
stepp = 0
for epoch in range(5):
    run_loss = 0.0

    for i,data in enumerate(train_loader,0):
        images,labels = data
        images,labels = images.cuda(),labels.cuda()
        images,labels = Variable(images),Variable(labels)

        out = net(images)
        loss = loss_fun(out,labels)

        opti.zero_grad()
        loss.backward()
        opti.step()
        print(i)
        print(loss.item())
        run_loss += loss.item()
        if i % 10 == 9:
            stepp += 10
            writer.add_scalar('train_loss',loss.item(),stepp)
        if i%2000 == 1999:
            print('[%d %d]: %.4f'%(epoch+1,i+1,run_loss))
            run_loss = 0.0

print('finish train')

torch.save(net.state_dict(),r'D:\Python\project\venv\Include\Nural_net/torch_net0')
print('save finish')


net.eval()

correct = 0.0
total = 0

for data in test_loader:
    images, labels = data
    images, labels = images.cuda(), labels.cuda()
    images, labels = Variable(images), Variable(labels)

    out = net(images)

    _,prediction = torch.max(out.data,1)
    print('prediction  ',prediction)
    print('labels  ',labels)
    total += labels.size(0)
    correct += ((prediction == labels).data).sum()
print('the acc is %.4f'%(correct*100/total))