pytorch载入预训练模型后,实现训练指定层

在PyTorch中，如果要载入预训练模型并对指定层进行训练，可以按照以下步骤进行操作：

1. 载入预训练模型
   在PyTorch中，载入预训练模型可以使用torchvision.models模块中的预置模型，例如resnet18。此外，如果需要使用自己的预训练模型，也可以使用torch.load()方法将之前训练好的模型载入。代码如下：

import torch
import torchvision.models as models

# 载入预置模型resnet18
model = models.resnet18(pretrained=True)

# 载入自己训练好的模型，假设模型保存在model.pth文件中
model = torch.load('model.pth')

1. 选定需要训练的层
   默认情况下，载入的预训练模型的所有层都是可以训练的。如果需要对指定层进行训练，可以先将所有层都设置为不可训练状态，然后将需要训练的层设置为可训练状态。代码如下：

for param in model.parameters():
    param.requires_grad = False   # 将所有层都设置为不可训练状态

# 设定需要训练的层
model.layer4[0].conv1.weight.requires_grad = True
model.layer4[0].bn1.weight.requires_grad = True
model.layer4[0].conv2.weight.requires_grad = True
model.layer4[0].bn2.weight.requires_grad = True

在上述代码中，我们将所有层都设置为不可训练状态，然后将layer4中的第一个卷积层、BatchNorm层、第二个卷积层和BatchNorm层设为可训练状态。

1. 进行训练和优化
   完成上述准备工作后，就可以进行模型训练和优化了。具体的训练和优化方法可以根据具体的需求而定，例如使用torch.optim.Adam优化器和交叉熵损失函数。这里不再赘述。

下面给出一个使用预置模型resnet18进行fine-tune的例子：

import torch
import torchvision.models as models
import torch.nn as nn

# 载入预置模型resnet18
model = models.resnet18(pretrained=True)

# 将所有层都设为不可训练状态，将最后一层全连接层抽出来
for param in model.parameters():
    param.requires_grad = False
fc_inputs = model.fc.in_features
model.fc = nn.Linear(fc_inputs, 2)

# 将最后一层的参数设为可训练状态
for param in model.fc.parameters():
    param.requires_grad = True

# 进行训练和优化
optimizer = torch.optim.Adam(model.fc.parameters())
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
    # 省略数据加载和前向传播部分

    loss = criterion(output, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

在上述代码中，我们将预置模型resnet18的所有层都设为不可训练状态，然后将最后一层的全连接层抽出来，设为可训练状态。最后使用交叉熵损失函数和Adam优化器进行训练。

下面给出一个使用自己的预训练模型进行fine-tune的例子：

import torch
import torch.nn as nn

# 载入自己训练好的模型，假设模型保存在model.pth文件中
model = torch.load('model.pth')

# 将所有层都设为不可训练状态，设定需要训练的层
for param in model.parameters():
    param.requires_grad = False   # 将所有层都设置为不可训练状态

model.layer4[0].conv1.weight.requires_grad = True
model.layer4[0].bn1.weight.requires_grad = True
model.layer4[0].conv2.weight.requires_grad = True
model.layer4[0].bn2.weight.requires_grad = True

# 进行训练和优化
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()))
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
    # 省略数据加载和前向传播部分

    loss = criterion(output, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

在上述代码中，我们先载入了自己训练好的模型，然后将所有层都设为不可训练状态，再将layer4中的第一个卷积层、BatchNorm层、第二个卷积层和BatchNorm层设为可训练状态。最后使用交叉熵损失函数和Adam优化器进行训练，注意优化器需要过滤掉不可训练的参数。