《深度学习之PyTorch实战计算机视觉》学习笔记（6）

这部分利用pytorch手动实现一个简化版的Vgg网络模型，用来对猫狗进行分类

由于CPU计算太慢，改用GPU来计算。代码基于python3.7， pytorch 1.0，cuda 10.0 .

import torch
import torchvision
from torch.autograd import Variable
from torchvision import datasets, transforms
import os            # 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)