图像分类（2），训练和预测_一般做分类预测需要多少训练集-优快云博客

本文链接：https://blog.youkuaiyun.com/u012327119/article/details/85073335

本文详细介绍了如何利用PyTorch的预训练模型进行微调，以适应特定的图像分类任务。覆盖了常见模型如ResNet、VGG等的加载与修改，以及损失函数、优化器的选择和训练流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

上篇博文分享了数据的预处理和封装，本文阐述一下如何用预训练模型进行训练和finetune。

pytorch已经准备了一些常用的网络模型，如alexnet、densenet、inception、resnet、squeezenet和vgg等。这些模型的原始定义文件可能在这个目录下/usr/lib/python2.7/site-packages/torchvision/models/。打卡resnet.py，可以看到：

__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
           'resnet152']

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

pytorch中定了已5种模型：resnet18， resnet34， resnet50， resnet101， resnet152，以及各自在imagenet上训练好的图像分类模型的网址。如果第一次运行以下代码，pytorch会从相应网址上下载预训练模型文件，保存在/root/.torch/models/中，后面再运行这两行代码，就直接从这个目录中载入模型，不需要下载了。

from torchvision import models
resnet18_pretrained = models.resnet18(pretrained=True)

可以看到，pytorch非常贴心的给用户准备了各种网络结构，如：

Alexnet： alexnet
DenseNet： densenet121， densenet169， densenet201， densenet161
Inception： Inception_v3
Resnet： resnet18, resnet34, resnet50, resnet101, resnet152
SqueezeNet: squeezenet1_0, squeezenet_1
VGG: vgg11,vgg13,vgg16,vgg19,vgg11_bn,vgg13_bn,vgg16_bn,vgg19_bn

但是这些预训练好的网络最后的全连接层的输出都是1000（用来进行1000个分类的），所以用这些预训练模型来训练自己的图像分类器，即需要更改全连接层的连接数，把1000改成自己数据集的类别数。更改也很简单，如下即可：

num_ftrs = resnet18_pretrained.fc.in_features
resnet18_pretrained.fc = nn.Linear(num_ftrs, 2)

然后，我们定义损失函数为交叉熵损失函数，这个损失函数度量了预测值和标签值的差距，我们的优化目标就是让神经网络的预测值竟可能接近与标签值，即最小化这个差距函数。如下

criterion = nn.CrossEntropyLoss()

pytorch0.4中定义了17中损失函数，如L1Loss, MSELoss, CrossEntropyLoss, NLLLoss, PossionNLLLoss, KLDivLoss, BCELoss, BCEWithLogitsLoss, MarginRankingLoss, MarginRankingLoss, HingeEmbeddingLoss, MultiLabelMarginLoss, SmoothL1Loss, SoftmarginLoss， MultiLabelSoftMarginLoss, CosineEmbeddingLoss, MultiMarginLoss, TripletMarginLoss等等。这些损失函数以后介绍。

有了损失函数，就要定义优化函数。优化函数是在torch.optim包中，包含了各种优化算法。优化方法包括optim.Adadelta, Adagrad, Adam, SparseAdam, Adamax, ASGD, LBFGS, RMSprop, Rprop, SGD等。

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

除了优化函数，还可以更改learning rate，在optim.lr_scheduler下有LambdaLR,StepLR, MultiStepLR, ExponentialLR, CosineAnnealingLR, ReduceLROnPlateau等，这里使用StepLR，表示每隔step_size个epoch，就将学习率将为原来的gamma倍。

接下来的训练就比较简单了，每一个batch训练数据，就先更改学习率值，把梯度值清零，然后计算预测值和标签值的loss，将loss梯度反向传播后，更新参数值。具体如下：

scheduler.step()
optimizer.zero_grad()
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()