利用pytorch实现迁移学习之猫狗分类器(dog vs cat)

最新推荐文章于 2025-07-04 15:24:50 发布

原创

最新推荐文章于 2025-07-04 15:24:50 发布 · 1.3w 阅读

68 ·

CC 4.0 BY-SA版权

迁移学习

迁移学习(Transfer learning) 就是把已学训练好的模型参数迁移到新的模型来帮助新模型训练。考虑到大部分数据或任务是存在相关性的，所以通过迁移学习我们可以将已经学到的模型参数（也可理解为模型学到的知识）通过某种方式来分享给新模型从而加快并优化模型的学习效率不用像大多数网络那样从零学习。
本文使用VGG16模型用于迁移学习，最终得到一个能对猫狗图片进行辨识的CNN（卷积神经网络），测试集用来验证模型是否能够很好的工作。

猫狗分类器

本文使用迁移学习实现猫狗分类器。
数据集来自Kaggle比赛：Dogs vs. Cats Redux: Kernels Edition

利用pytorch实现迁移学习

首先进行图片的导入和预览

path = "dog_vs_cat"
transform = transforms.Compose([transforms.CenterCrop(224),
                                transforms.ToTensor(),
                                transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])])

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"]}

输入的图片需要分辨率为224*224，所以使用transform.CenterCrop(224)对原始图片进行裁剪。载入的图片训练集合为20000个和验证集合为5000个，原始图片全部为训练集合，需自己拆分出一部分验证集合，输出的Label，1代表为狗，0代表猫。

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)