pytorch实现迁移学习之猫狗大战Demo
参考链接：https://blog.csdn.net/Sugar_girl/article/details/80050690
注意事项：
操作系统：Windows7 （在Linux下更好，至少不用担心路径的转义字符问题了）
Python：Python3.6 （Python3都可以，最好用Anaconda中的Python解释器，里边预安装了很多包）
数据集：数据集最好不要 kaggle 的猫狗全集，不然时间花费太长了...
目录结构为：
└─kaggle_dog_vs_cat
├─pretrain
├─test
│ └─test
├─train
│ ├─cat
│ └─dog
├─val
│ ├─cat
│ └─dog
└─zip其中目录 pretrain 中主要有一个下载的 vgg16-397923af.pth 文件； .\test\test 目录下是你最后要进行分类测试的猫狗图片；zip文件夹下是解压之后所有的猫狗图片；train 目录下是从 zip 目录下分离出的 90% 的猫狗图片作为训练集，且猫狗图片在各自的子目录下；val 目录下是从 zip 目录下分离出的 10% 的猫狗图片作为验证集，且猫狗图片在各自的子目录下；
预处理：预处理主要是把所有训练集分成训练集和验证集，训练集占的比例一般会高一些，但各自所占的比例你可以自己决定。
"""
@author: Looking
@email: [email protected]
"""
import os
import shutil
data_file = os.listdir("D:\MyProject\Python\MachineLearning\deeplearning\kaggle_dog_vs_cat\zip")
dog_file = list(filter(lambda x: x[:3] == 'dog', data_file))
cat_file = list(filter(lambda x: x[:3] == 'cat', data_file))
root = r"D:\MyProject\Python\MachineLearning\deeplearning\kaggle_dog_vs_cat"
train_root = root + r"\train"
val_root = root + r"\val"
print(root)
print(train_root)
print(val_root)
print("begin data remove.")
for i in range(len(dog_file)):
pic_path = root + 'zip\\' + dog_file[i]
if i < len(dog_file) * 0.9:
obj_path = train_root + '\\dog\\' + dog_file[i]
else:
obj_path = val_root + '\\dog\\' + dog_file[i]
shutil.move(pic_path, obj_path)
for i in range(len(cat_file)):
pic_path = root + 'zip\\' + cat_file[i]
if i < len(cat_file) * 0.9:
obj_path = train_root + '\\cat\\' + cat_file[i]
else:
obj_path = val_root + '\\cat\\' + cat_file[i]
shutil.move(pic_path, obj_path)
print("data remove finished.")猫狗分类器Demo：
"""
@author: Looking
@email: [email protected]
"""
import torch
import torchvision
from torchvision import datasets, transforms, models
import os
import numpy as np
import matplotlib.pyplot as plt
from torch.autograd import Variable
import time
transform = transforms.Compose([transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
path = "D:\MyProject\Python\MachineLearning\deeplearning\kaggle_dog_vs_cat"
data_image = {x: datasets.ImageFolder(root=os.path.join(path, x), transform=transform)
for x in ["train", "val"]}
data_loader_image = {x: torch.utils.data.DataLoader(dataset=data_image[x], batch_size=4,
shuffle=True) for x in ["train", "val"]}
print(data_image)
use_gpu = torch.cuda.is_available()
classes = data_image["train"].classes
# classes_index = data_image["train"].class_to_idx
# print("train data set:", len(data_image["train"]))
# print("val data set:", len(data_image["val"]))
# X_train, y_train = next(iter(data_loader_image["train"]))
# mean = [0.5, 0.5, 0.5]
# std = [0.5, 0.5, 0.5]
# img = torchvision.utils.make_grid(X_train)
# img = img.numpy().transpose((1, 2, 0))
# img = img * std + mean
#
# print([classes[i] for i in y_train])
# plt.imshow(img)
# plt.show()
#
#
#
model = models.vgg16(pretrained=False)
pre = torch.load(r'.\kaggle_dog_vs_cat\pretrain\vgg16-397923af.pth')
model.load_state_dict(pre)
# model = models.vgg16(pretrained=True) # 也可以直接从官网的url加载，但是很慢，经常出现连接超时错误。
# print(model)
for parma in model.parameters():
parma.requires_grad = False
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 4096),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(4096, 4096),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(4096, 2))
for index, parma in enumerate(model.classifier.parameters()):
if index == 6:
parma.requires_grad = True
use_gpu = False # 我用GPU加速会出现 RuntimeError: CUDA out of memory错误，所以没用GPU。
if use_gpu:
model = model.cuda()
cost = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.classifier.parameters())
print(model)
n_epochs = 1
for epoch in range(n_epochs):
since = time.time()
print("Epoch{}/{}".format(epoch, n_epochs))
print("-" * 10)
for param in ["train", "val"]:
if param == "train":
model.train = True
else:
model.train = False
running_loss = 0.0
running_correct = 0
batch = 0
for data in data_loader_image[param]:
batch += 1
X, y = data
if use_gpu:
X, y = Variable(X.cuda()), Variable(y.cuda())
else:
X, y = Variable(X), Variable(y)
optimizer.zero_grad()
y_pred = model(X)
_, pred = torch.max(y_pred.data, 1)
loss = cost(y_pred, y)
if param == "train":
loss.backward()
optimizer.step()
running_loss += loss.item()
running_correct += torch.sum(pred == y.data)
if batch % 5 == 0 and param == "train":
print("Batch {}, Train Loss:{:.4f}, Train ACC:{:.4f}".format(
batch, running_loss / (4 * batch), 100 * running_correct / (4 * batch)))
epoch_loss = running_loss / len(data_image[param])
epoch_correct = 100 * running_correct / len(data_image[param])
print("{} Loss:{:.4f}, Correct:{:.4f}".format(param, epoch_loss, epoch_correct))
now_time = time.time() - since
print("Training time is:{:.0f}m {:.0f}s".format(now_time // 60, now_time % 60))
test_path = r"D:\MyProject\Python\MachineLearning\deeplearning\kaggle_dog_vs_cat\test\\"
# 目录结构中虽然是两个test，但是这儿的test_path路径中只有一个test，切记切记
data_test_img = datasets.ImageFolder(root=test_path, transform=transform)
data_loader_test_img = torch.utils.data.DataLoader(dataset=data_test_img,
batch_size=16)
image, label = next(iter(data_loader_test_img))
images = Variable(image)
y_pred = model(images)
_, pred = torch.max(y_pred.data, 1)
print(pred)
img = torchvision.utils.make_grid(image)
img = img.numpy().transpose(1, 2, 0)
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
img = img * std + mean
print("Pred Label:", [classes[i] for i in pred])
plt.imshow(img)
plt.show()
model = model.load_state_dict(pre)
print(model)
如果程序中的路径都是使用的绝对路径，那么 .py 文件所在的位置已经不是很重要了，不过最好可以放在 kaggle_dog_vs_cat 的同级目录。