使用深度学习进行图像分类pytorch

from glob import glob
import os
import numpy as np
import matplotlib.pyplot as plt
import shutil
from torchvision import transforms
from torchvision import models
import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.optim import lr_scheduler
from torch import optim
from torchvision.datasets import ImageFolder
from torchvision.utils import make_grid
import time

def imshow(inp):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)

path = '../Chapter03/dogsandcats/'
files = glob(os.path.join(path,'*/*.jpg'))#读取文件夹内的所有文件，join连接两个或更多的路径名组件
no_of_images = len(files)
shuffle = np.random.permutation(no_of_images)#创建可用于创建验证数据集的混合索引
os.mkdir(os.path.join(path,'valid'))#创建保存验证图片集的 validation 目录
os.mkdir(os.path.join(path,'train'))
for t in ['train','valid']:#使用标签名称创建目录
    for folder in ['dog/','cat/']:
        os.mkdir(os.path.join(path,t,folder)) 

for i in shuffle[:2000]:#将图片的一小部分子集复制到 validation 文件夹
    #shutil.copy(files[i],'../Chapter03/dogsandcats/valid/')
#     folder = files[i].split('/')[-1].split('.')[0]
#     image = files[i].split('/')[-1]
#     os.rename(files[i],os.path.join(path,'valid',folder,image))
    if files[i].split('\\')[-1].split('.')[0]=="cat":
        shutil.copy(files[i],'../Chapter03/dogsandcats/valid/cat/')
    else:
        shutil.copy(files[i],'../Chapter03/dogsandcats/valid/dog/')

for i in shuffle[2000:]:#将图片的一小部分子集复制到 training 文件夹
    #shutil.copy(files[i],'../chapter3/dogsandcats/train/')
#     folder = files[i].split('/')[-1].split('.')[0]
#     image = files[i].split('/')[-1]
#     os.rename(files[i],os.path.join(path,'train',folder,image))
    if files[i].split('\\')[-1].split('.')[0]=="cat":
        shutil.copy(files[i],'../Chapter03/dogsandcats/train/cat/')
    else:
        shutil.copy(files[i],'../Chapter03/dogsandcats/train/dog/')

if torch.cuda.is_available():
    is_cuda = True

simple_transform = transforms.Compose([transforms.Resize((224,224))
                                       ,transforms.ToTensor()
                                       ,transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
train = ImageFolder('dogsandcats/train/',simple_transform)
valid = ImageFolder('dogsandcats/valid/',simple_transform)

train_data_gen = torch.utils.data.DataLoader(train,shuffle=True,batch_size=64,num_workers=3)
valid_data_gen = torch.utils.data.DataLoader(valid,batch_size=64,num_workers=3)

dataset_sizes = {'train':len(train_data_gen.dataset),'valid':len(valid_data_gen.dataset)}
dataloaders = {'train':train_data_gen,'valid':valid_data_gen}

model_ft = models.resnet18(pretrained=True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 2)#把输出特征改成 2

if torch.cuda.is_available():
    model_ft = model_ft.cuda()

# Loss and Optimizer
learning_rate = 0.001
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)#StepLR 函数帮助动态修改学习率

def train_model(model, criterion, optimizer, scheduler, num_epochs=5): 
    since = time.time()
    best_model_wts = model.state_dict()
#    pytorch 中的 state_dict 是一个简单的python的字典对象,将每一层与它的对应参数建立映射关系.(如model的每一层的weights及偏置等等)
#    (注意,只有那些参数可以训练的layer才会被保存到模型的state_dict中,如卷积层,线性层等等)
#    优化器对象Optimizer也有一个state_dict,它包含了优化器的状态以及被使用的超参数(如lr, momentum,weight_decay等)
    best_acc = 0.0
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
        #每轮都有训练和验证阶段
        for phase in ['train', 'valid']:
            if phase == 'train':
                scheduler.step()
                model.train(True)  # 模型设为训练模式
            else:
                model.train(False)  # 模型设为评估模式
            running_loss = 0.0
            running_corrects = 0
            #在数据上迭代
            for data in dataloaders[phase]:
                # 获取输入
                inputs, labels = data
                # 封装成变量
                if torch.cuda.is_available():
                    inputs = Variable(inputs.cuda())
                    labels = Variable(labels.cuda())
                else:
                    inputs, labels = Variable(inputs), Variable(labels)
                #梯度参数清 0
                optimizer.zero_grad()
                #前向
                outputs = model(inputs)
                _, preds = torch.max(outputs.data, 1)#torch.max(a,1) 返回每一行中最大值的那个元素，且返回其索引（返回最大元素在这一行的列索引）
                loss = criterion(outputs, labels)
                #只在训练阶段反向优化
                if phase == 'train':
                    loss.backward()
                    optimizer.step()
                #统计
                running_loss += loss.data
                running_corrects += torch.sum(preds == labels.data)
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects / dataset_sizes[phase]
            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
            #深度复制模型
            if phase == 'valid' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = model.state_dict()
        print()
    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
    #加载最优权重
    model.load_state_dict(best_model_wts)
    return model

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,num_epochs=5)

以上摘自《PyTorch深度学习》书籍。