实战人工智能框架PyTorch的模型训练

我们关心网络结构和数据，定义损失函数，定义优化函数等。

具体步骤如下：

第一步：图像封装为向量后，将输入input向前传播，进行运算后得到输出output

第二步：将output再输入loss函数，计算loss值（是个标量）-损失函数用来得到新权重

第三步：将梯度反向传播到每个参数（优化函数）：主要指标是 学习速率?  x 梯度向量g

第四步：利用下面公式进行权重更新

新权重w =  旧权重w  +  学习速率?  x 梯度向量g

封装好了数据后，就可以作为模型的输入了。所以要先导入你的模型。在PyTorch中已经默认为大家准备了一些常用的网络结构，比如分类中的VGG，ResNet，DenseNet等等，可以用torchvision.models模块来导入。比如用torchvision.models.resnet18(pretrained=True)来导入ResNet18网络，同时指明导入的是已经预训练过的网络。因为预训练网络一般是在1000类的ImageNet数据集上进行的，所以要迁移到你自己数据集的2分类，需要替换最后的全连接层为你所需要的输出。因此下面这三行代码进行的就是用models模块导入resnet18网络，然后获取全连接层的输入channel个数，用这个channel个数和你要做的分类类别数（这里是2）替换原来模型中的全连接层。这样网络结果也准备好。

model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

但是只有网络结构和数据还不足以让代码运行起来，还需要定义损失函数。在PyTorch中采用torch.nn模块来定义网络的所有层，比如卷积、降采样、损失层等等，这里采用交叉熵函数，因此可以这样定义：

criterion = nn.CrossEntropyLoss()
然后你还需要定义优化函数，比如最常见的随机梯度下降，在PyTorch中是通过torch.optim模块来实现的。另外这里虽然写的是SGD，但是因为有momentum，所以是Adam的优化方式。这个类的输入包括需要优化的参数：model.parameters()，学习率，还有Adam相关的momentum参数。现在很多优化方式的默认定义形式就是这样的。

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

然后一般还会定义学习率的变化策略，这里采用的是torch.optim.lr_scheduler模块的StepLR类，表示每隔step_size个epoch就将学习率降为原来的gamma倍。

代码实现如下:

from __future__ import print_function, division

1.导入算法的包

import torch

import torch.nn as nn

import torch.optim as optim

from torch.optim import lr_scheduler

from torch.autograd import Variable

import torchvision

from torchvision import datasets, models, transforms

import time

import os

2.训练模型

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):

since = time.time()

best_model_wts = model.state_dict()

best_acc = 0.0

for epoch in range(num_epochs):

print('Epoch {}/{}'.format(epoch, num_epochs - 1))

print('-' * 10)

# Each epoch has a training and validation phase

for phase in ['train', 'val']:

if phase == 'train':

scheduler.step()

model.train(True)  # Set model to training mode

else:

model.train(False)  # Set model to evaluate mode

running_loss = 0.0

running_corrects = 0.0

# Iterate over data.

for data in dataloders[phase]:

# get the inputs

inputs, labels = data

# wrap them in Variable

if use_gpu:

inputs = Variable(inputs.cuda())

labels = Variable(labels.cuda())

else:

inputs, labels = Variable(inputs), Variable(labels)

# zero the parameter gradients

optimizer.zero_grad()

# forward

outputs = model(inputs)

_, preds = torch.max(outputs.data, 1)

loss = criterion(outputs, labels)

# backward + optimize only if in training phase

if phase == 'train':

loss.backward()

optimizer.step()

# statistics

running_loss += loss.data[0]

running_corrects += torch.sum(preds == labels.data).to(torch.float32)

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))

# deep copy the model

if phase == 'val' and epoch_acc > best_acc:

best_acc = epoch_acc

best_model_wts = model.state_dict()

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))

# load best model weights

model.load_state_dict(best_model_wts)

return model

if __name__ == '__main__':

# data_transform, pay attention that the input of Normalize() is Tensor and the input of RandomResizedCrop() or RandomHorizontalFlip() is PIL Image

data_transforms = {

'train': transforms.Compose([

transforms.RandomSizedCrop(224),

transforms.RandomHorizontalFlip(),

transforms.ToTensor(),

transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

]),

'val': transforms.Compose([

transforms.Scale(256),

transforms.CenterCrop(224),

transforms.ToTensor(),

transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

]),

}

转换为元组image_datasets

# your image data file

data_dir = '/data'

image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),

data_transforms[x]) for x in ['train', 'val']}

# wrap your data and label into Tensor 把image_datasets元组转换为Tensor

dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],

batch_size=4,

shuffle=True,

num_workers=4) for x in ['train', 'val']}

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}

# use gpu or not

use_gpu = torch.cuda.is_available()

# get model and replace the original fc layer with your fc layer

model_ft = models.resnet18(pretrained=True)

num_ftrs = model_ft.fc.in_features

model_ft.fc = nn.Linear(num_ftrs, 2)

if use_gpu:

model_ft = model_ft.cuda()

# define loss function

criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized

optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

model_ft = train_model(model=model_ft,

criterion=criterion,

optimizer=optimizer_ft,

scheduler=exp_lr_scheduler,

num_epochs=25)