pytorch 的torchvision中继承了常用的数据集，可以直接在torchvision.datasets中调用，当我们想要加载自己的数据集来训练模型时，需要借助torchvision.datasets.ImageFolder这个类，在当前的Pytorch版本中， ImageFolder继承自DatasetFolder官方源码
我们可以直接使用ImageFolder进行数据的读取，但是这要求我们的数据集按照以下方式被组织：
''' A generic data loader where the samples are arranged in this way: :: root/class_x/xxx.ext root/class_x/xxy.ext root/class_x/xxz.ext root/class_y/123.ext root/class_y/nsdf3.ext root/class_y/asd932_.ext '''即保证在文件根目录root下，每一个类别的图像被置于一个以标签为名的文件夹中。
import torch import torch.nn as nn import torchvision.datasets as datasets import torchvision.transforms as transforms import torchvision.models as models import torch.nn.functional as F import torch.optim as optim import os import time from datetime import timedelta #指定训练数据与测试数据的路径 traindir = r'F:\DeepLearning\datasets\dogvscat\training_set' testdir = r'F:\DeepLearning\datasets\dogvscat\test_set' batch_size = 32 num_worker = 2 #定义加载数据函数load_data() 返回train_loder和test_loder def load_data(traindir, testdir, batch_size = batch_size, num_worker = num_worker, shuffle = True): print('starting load data...') train_loder = torch.utils.data.DataLoader( datasets.ImageFolder(traindir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor() ])), batch_size= batch_size, shuffle = shuffle, num_workers = num_worker) test_loder = torch.utils.data.DataLoader( datasets.ImageFolder(testdir, transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor() ])), batch_size= batch_size, shuffle = shuffle, num_workers = num_worker) print('ending load data...') return train_loder, test_loder 在进行神经网络模型的训练时，我们往往不会重头开始训练一个网络，而是借助一些在大型数据集上已经训练好的模型完成特征的提取，再增加一些自定义的层完成深度学习任务（如：分类）。在训练的过程中需要冻结预训练模型的各层，只训练自定义的层。
在这个Dogs vs Cats的分类任务中，我主要使用了预训练的VGG16模型，可以直接在torchvision.models中使用vgg16, 并在最后一层FC后增加了一层线性层nn.Linear(input_features=1000, output_features=2)，自定义模型Mymodel:
class Mymodel(nn.Module): def __init__(self): super(Mymodel, self).__init__() self.vgg = models.vgg16() self.last_layer = torch.nn.Linear(in_features=1000, out_features=2) def forward(self, x): x = self.vgg(x) x = F.relu(x) x = self.last_layer(x) return x随后build_model()函数完成预训练模型的权重加载与冻结：
def build_model(path): #加载vgg预训练参数 vgg_state_dict = torch.load(path) model = Mymodel() model.vgg.load_state_dict(vgg_state_dict) #冻结部分参数 for name, p in model.vgg.named_parameters(): p.requires_grad = False return model在定义模型之后可以查看以下模型结构：
print(model) Mymodel( (vgg): VGG( (features): Sequential( (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): ReLU(inplace=True) (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (3): ReLU(inplace=True) (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (6): ReLU(inplace=True) (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (8): ReLU(inplace=True) (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (11): ReLU(inplace=True) (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (13): ReLU(inplace=True) (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (15): ReLU(inplace=True) (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (18): ReLU(inplace=True) (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (20): ReLU(inplace=True) (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (22): ReLU(inplace=True) (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (25): ReLU(inplace=True) (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (27): ReLU(inplace=True) (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (29): ReLU(inplace=True) (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) ) (avgpool): AdaptiveAvgPool2d(output_size=(7, 7)) (classifier): Sequential( (0): Linear(in_features=25088, out_features=4096, bias=True) (1): ReLU(inplace=True) (2): Dropout(p=0.5, inplace=False) (3): Linear(in_features=4096, out_features=4096, bias=True) (4): ReLU(inplace=True) (5): Dropout(p=0.5, inplace=False) (6): Linear(in_features=4096, out_features=1000, bias=True) ) ) (last_layer): Linear(in_features=1000, out_features=2, bias=True) )定义训练函数train(), 参数包括：
data_loder: 数据集model: 需要训练的模型lossFunc： 损失函数（eg:torch.nn.CrossEntropyLoss())optimizer: 优化器（eg: torch.optim.Adam) def train(data_loder, model, lossFunc, optimizer, num_epoch = 1): print('Start training...') for epoch in range(num_epoch): print('Starting epoch %d / %d' %(epoch+1, num_epoch)) #模型训练 for i,(images, labels) in enumerate(data_loder): predict = model(images) optimizer.zero_grad() loss = lossFunc(predict, labels) loss.backward() optimizer.step() print('epoch %d, iteration %d, loss %s' % (epoch, i, loss.item())) print('End training...')先可视化部分测试数据：
import matplotlib.pyplot as plt batch_images = batch_images.permute(0,2,3,1) #对于多维tensor使用permute() f = plt.figure(figsize=(12,6)) for i in range(4): f.add_subplot(1,4, i+1) plt.imshow(batch_images[i]) #在测试集上测试准确性 #在测试集上测试性能 total = 0 correct = 0 iter_data = iter(test_loder) with torch.no_grad(): for i in range(20): data = iter_data.next() images, labels = data out = model(images) _, predict = torch.max(out, 1) total += labels.size(0) correct += (labels == predict).sum().item() print('Accuracy: %s' % (100*correct / total)) Accuracy: 94.21875由于算力的问题，无法训练多个epoch， 并进行模型的fine-tuning，准确率有点小低。。。
Pytorch中有两种保存模型的方式，一种是以state_dict的形式保存所有的模型参数，模型恢复时可以调用load_state_dict()进行恢复：
#保存模型参数 torch.save(model.state_dict(), path) #恢复模型 model = Mymodel() state_dict = torch.load(path) model.load_state_dict(state_dict)只保存模型参数的方式执行起来比较快。
另一种保存模型的方法是保存模型结构+参数：
#保存完整模型 torch.save(model, path) #恢复完整模型 model = torch.load(path)这个例子中我采用了第一种方式：
#模型保存参数 save_dir = r'F:\DeepLearning\models' torch.save(model.state_dict(), os.path.join(save_dir, 'mymodel.pkl'))