我们关心网络结构和数据,定义损失函数,定义优化函数等。
具体步骤如下:
第一步:图像封装为向量后,将输入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) |