pytorch猫狗大战-项目代码


各位看官老爷,小白我知道pytorch的这个猫狗大战的代码真是漫天都是,这基本上也就是一个hello world的程序。我这个代码也是从《pytorch 机器学习从入门到实践》一书中摘出来了。对于初学者来说(像小白我)这本书还是不错的,难能可贵的是,它良心的在github上给出了代码。如果各位看官也有这本书的话,一定不要直接敲书中的代码,而是要结合着github上的代码再敲。因为这本书中的印刷的代码有缩进错误和一些疏漏的地方。有时照着这个书中的代码敲是会报错的。书中的代码只是一个核心代码的参考。我这也是敲过之后的感悟,现在就直接贴出这个hello word代码,本人亲测是可以运行的通的。感兴趣的大神就运行一下吧。


0. 前提:

已经下载了kaggle中的猫狗数据集

1. 数据预处理代码

将下载好的数据集中的trainzip处理成一个小的数据集,并且区分出train和test比例为0.95和0.05。并将其中的猫狗图像进行归一化处理

import os, shutil
import numpy as np 
import pdb

random_state = 42
np.random.seed(random_state)

original_dataset_dir = './data/CAT-DOG/train' #这是我的train的路径可以换成你的
total_num = int(len(os.listdir(original_dataset_dir)) / 2)
random_idx = np.array(range(total_num))
np.random.shuffle(random_idx)
base_dir = 'cats_and_dogs_small'
if not os.path.exists(base_dir):
	os.mkdir(base_dir)

sub_dirs = ['train', 'test']
animals = ['cats', 'dogs']
train_idx = random_idx[: int(total_num * 0.9)]
test_idx = random_idx[int(total_num * 0.9) :]
numbers = [train_idx, test_idx]
for idx, sub_dir in enumerate(sub_dirs):
	dir = os.path.join(base_dir, sub_dir)
	if not os.path.exists(dir):
		os.mkdir(dir)
	for animal in animals:
		animal_dir = os.path.join(dir, animal)
		if not os.path.exists(animal_dir):
			os.mkdir(animal_dir)
		fnames = [animal[: -1]+ '.{}.jpg'.format(i) for i in numbers[idx]]
		for fname in fnames:
			src = os.path.join(original_dataset_dir, fname)
			dst = os.path.join(animal_dir,fname)
			shutil.copyfile(src, dst)
		print(dir + 'total images : %d'%(len(os.listdir(animal_dir))))

 2. 自己设计的网络中的训练和测试代码

这部分代码在书中是存在部分问题的,但是我这是修改后的能够运行的代码

from __future__ import print_function, division
import shutil
import torch
import os
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np 
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, datasets, utils
from torch.utils.data import DataLoader
import torch.optim as optim

torch.manual_seed(1)
epochs = 10
batch_size = 4
num_workers = 4
use_gpu = torch.cuda.is_available()

data_transform = transforms.Compose([transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])])


train_dataset = datasets.ImageFolder(root = 'cats_and_dogs_small/train/', transform = data_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = batch_size, shuffle = True, num_workers = num_workers)
test_dataset = datasets.ImageFolder(root = 'cats_and_dogs_small/test/', transform=data_transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = batch_size, shuffle=True, num_workers = num_workers)

class Net(nn.Module):
	def __init__(self):
		super(Net, self).__init__()
		self.conv1 = nn.Conv2d(3, 6, 5)
		self.maxpool = nn.MaxPool2d(2, 2)
		self.conv2 = nn.Conv2d(6, 16, 5)
		self.fc1 = nn.Linear(16 * 53 * 53, 1024)
		self.fc2 = nn.Linear(1024, 512)
		self.fc3 = nn.Linear(512, 2)
	def forward(self, x):
		x = self.maxpool(F.relu(self.conv1(x)))
		x = self.maxpool(F.relu(self.conv2(x)))
		x = x.view(-1, 16 * 53 * 53)
		x = F.relu(self.fc1(x))
		x = F.relu(self.fc2(x))
		x = self.fc3(x)
		return x
if use_gpu:
	net = Net().cuda()
else:
	net = Net()
print(net)

cirterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.0001, momentum=0.9)

net.train()
for epoch in range(epochs):
	running_loss = 0.0
	train_correct = 0
	train_total = 0
	for i, data in enumerate(train_loader, 0):
		inputs, train_labels = data
		if use_gpu:
			inputs, labels = Variable(inputs.cuda()), Variable(train_labels.cuda())
		else:
			inputs, labels = Variable(inputs), Variable(train_labels)

		optimizer.zero_grad()
		outputs = net(inputs)
		_, train_predicted = torch.max(outputs.data, 1)

		train_correct += (train_predicted == labels.data).sum()
		loss = cirterion(outputs, labels)
		loss.backward()
		optimizer.step()
		running_loss += loss.item()
		train_total += train_labels.size(0)
	print('train %d epoch loss: %.3f acc: %.3f '% (epoch + 1, running_loss / train_total, 100 * train_correct / train_total))

	correct = 0
	test_loss = 0.0
	test_total = 0
	net.eval()
	for data in test_loader:
		images, labels = data
		if use_gpu:
			images, labels = Variable(images.cuda()), Variable(labels.cuda())
		else:
			images, labels = Variable(images), Variable(labels)
		outputs = net(images)
		_, predicted = torch.max(outputs.data, 1)
		loss = cirterion(outputs, labels)
		test_loss += loss.item()
		test_total += labels.size(0)
		correct += (predicted == labels.data).sum()

	print('test %d epoch loss: %.3f acc: %.3f' % (epoch + 1, test_loss / test_total, 100 * correct / test_total))

2. 迁移学习用resnet网络进行训练和预测的代码

# coding=utf-8

from __future__ import print_function, division
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import Dataset
from torchvision import transforms, datasets, models


# 配置参数
random_state = 1
torch.manual_seed(random_state)  # 设置随机数种子,确保结果可重复
torch.cuda.manual_seed(random_state)
torch.cuda.manual_seed_all(random_state)
np.random.seed(random_state)
# random.seed(random_state)

epochs = 10  # 训练次数
batch_size = 4  # 批处理大小
num_workers = 4  # 多线程的数目
use_gpu = torch.cuda.is_available()
# 对加载的图像作归一化处理, 并裁剪为[224x224x3]大小的图像
data_transform = transforms.Compose([
    transforms.Scale(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 数据的批处理,尺寸大小为batch_size,
# 在训练集中,shuffle 必须设置为True, 表示次序是随机的
train_dataset = datasets.ImageFolder(root='cats_and_dogs_small/train/', transform=data_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)

test_dataset = datasets.ImageFolder(root='cats_and_dogs_small/test/', transform=data_transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)





# 创建模型

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.maxpool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 53 * 53, 1024)
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 2)

    def forward(self, x):
        x = self.maxpool(F.relu(self.conv1(x)))
        x = self.maxpool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 53 * 53)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 加载resnet18 模型,
net = models.resnet18(pretrained=False)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 2)  # 更新resnet18模型的fc模型,

if use_gpu:
    net = net.cuda()
print(net)
# 定义loss和optimizer
cirterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.0001, momentum=0.9)


# 开始训练
net.train()
for epoch in range(epochs):
    running_loss = 0.0
    train_correct = 0
    train_total = 0
    for i, data in enumerate(train_loader, 0):
        inputs, train_labels = data
        if use_gpu:
            inputs, labels = Variable(inputs.cuda()), Variable(train_labels.cuda())
        else:
            inputs, labels = Variable(inputs), Variable(train_labels)
        # inputs, labels = Variable(inputs), Variable(train_labels)
        optimizer.zero_grad()
        outputs = net(inputs)
        _, train_predicted = torch.max(outputs.data, 1)
        # import pdb
        # pdb.set_trace()
        train_correct += (train_predicted == labels.data).sum()
        loss = cirterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        train_total += train_labels.size(0)

    print('train %d epoch loss: %.3f  acc: %.3f ' % (
    epoch + 1, running_loss / train_total, 100 * train_correct / train_total))

    # 模型测试
    correct = 0
    test_loss = 0.0
    test_total = 0
    test_total = 0
    net.eval()
    for data in test_loader:
        images, labels = data
        if use_gpu:
            images, labels = Variable(images.cuda()), Variable(labels.cuda())
        else:
            images, labels = Variable(images), Variable(labels)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        loss = cirterion(outputs, labels)
        test_loss += loss.item()
        test_total += labels.size(0)
        correct += (predicted == labels.data).sum()
    print('test  %d epoch loss: %.3f  acc: %.3f ' % (epoch + 1, test_loss / test_total, 100 * correct / test_total))

 希望各位看官敲代码愉快!