当前位置：首页 > news >正文

pytorch学习笔记（五）-- 计算机视觉的迁移学习

news 2025/7/18 6:09:12

系列文章目录

pytorch学习笔记（一）-- pytorch深度学习框架基本知识了解

pytorch学习笔记（二）-- pytorch模型开发步骤详解

pytorch学习笔记（三）-- TensorBoard的介绍

pytorch学习笔记（四）-- TorchVision 物体检测微调教程

pytorch学习笔记（五）-- 计算机视觉的迁移学习

文章目录

系列文章目录

文章目录

前言

一、加载数据

二、训练模型

三、可视化模型预测

四、卷积网络微调

微调 ConvNet：

ConvNet 作为固定特征提取器：

五、自定义图像的推理

总结

前言

在本章节，您将学习如何使用迁移学习训练卷积神经网络进行图像分类。您可以在 cs231n notes 笔记中阅读有关迁移学习的更多信息。

一般来说，大家都不会从头开始训练卷积神经网络，而是先在较大的数据集上做预训练，差不多成熟了，然后再把卷积网络在自己的任务上做初始化，或者特征提取器。

这两种主要的迁移学习场景如下：

微调 ConvNet：我们不是使用随机初始化，而是使用预训练网络（例如在 imagenet 1000 数据集上训练的网络）来初始化网络。其余训练看起来与往常一样。
ConvNet 作为固定特征提取器：在这里，我们将冻结除最终全连接层之外的所有网络的权重。最后一个全连接层将被替换为具有随机权重的新层，并且只训练这一层。

一、加载数据

我们使用torchvision 和 torch.utils.data数据包进行数据加载，今天的任务是训练一个模型用来分辨蚂蚁和蜜蜂，我们有120张蚂蚁和蜜蜂的照片用于训练，以及75张用于测试蜜蜂和蚂蚁的照片。

数据集下载路径：MyDataset: 数据集仓库，包括各种网站搜刮的，以及一些自定义的数据。方便后续神经网络的训练 - Gitee.com

通常，如果从头开始训练，这个数据集太小了，无法进行推广。由于我们使用迁移学习，我们就可以相当好地进行推广。

# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {'train': transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),'val': transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
}data_dir = 'data/hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),data_transforms[x])for x in ['train', 'val']}
dataloaders = {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']}
class_names = image_datasets['train'].classesdevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#可视化部分图片，确认下效果
def imshow(inp, title=None):"""Display image 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 + meaninp = np.clip(inp, 0, 1)plt.imshow(inp)if title is not None:plt.title(title)plt.pause(0.001)  # pause a bit so that plots are updated# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))# Make a grid from batch
out = torchvision.utils.make_grid(inputs)
imshow(out, title=[class_names[x] for x in classes])

二、训练模型

编写一个通用函数来训练模型。

安排这个learning rate
保存模型

参数 scheduler 是来自 torch.optim.lr_scheduler 的 LR 调度程序对象，关于这个scheduler在后续模型优化的章节会讲到，也是一个非常强大的功能，这里就先不赘述了。

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):since = time.time()# Create a temporary directory to save training checkpointswith TemporaryDirectory() as tempdir:best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')torch.save(model.state_dict(), best_model_params_path)best_acc = 0.0for epoch in range(num_epochs):print(f'Epoch {epoch}/{num_epochs - 1}')print('-' * 10)# Each epoch has a training and validation phasefor phase in ['train', 'val']:if phase == 'train':model.train()  # Set model to training modeelse:model.eval()   # Set model to evaluate moderunning_loss = 0.0running_corrects = 0# Iterate over data.for inputs, labels in dataloaders[phase]:inputs = inputs.to(device)labels = labels.to(device)# zero the parameter gradientsoptimizer.zero_grad()# forward# track history if only in trainwith torch.set_grad_enabled(phase == 'train'):outputs = model(inputs)_, preds = torch.max(outputs, 1)loss = criterion(outputs, labels)# backward + optimize only if in training phaseif phase == 'train':loss.backward()optimizer.step()# statisticsrunning_loss += loss.item() * inputs.size(0)running_corrects += torch.sum(preds == labels.data)if phase == 'train':scheduler.step()epoch_loss = running_loss / dataset_sizes[phase]epoch_acc = running_corrects.double() / dataset_sizes[phase]print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')# deep copy the modelif phase == 'val' and epoch_acc > best_acc:best_acc = epoch_acctorch.save(model.state_dict(), best_model_params_path)print()time_elapsed = time.time() - sinceprint(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')print(f'Best val Acc: {best_acc:4f}')# load best model weightsmodel.load_state_dict(torch.load(best_model_params_path, weights_only=True))return model

三、可视化模型预测

经过上一节的训练，我们现在看看模型的预测结果怎么样。作为传统的程序开发者，我们习惯通过打印来验证结果，但是这玩意儿只有开发兄弟能懂哈。Pytorch毕竟是基于Python的深度学习框架，工具包是应用尽有，所以我们可以以可视化的图形来显示预测的效果。说个题外话，将工作结果可视化这个习惯，各位开发兄弟得学起来，以后就可以一手抓开发，一手抓产品，既可以跟开发同事一起奋斗，又可以跟老板以及客户吹牛逼，路就走宽了。

def visualize_model(model, num_images=6):was_training = model.trainingmodel.eval()images_so_far = 0fig = plt.figure()with torch.no_grad():for i, (inputs, labels) in enumerate(dataloaders['val']):inputs = inputs.to(device)labels = labels.to(device)outputs = model(inputs)_, preds = torch.max(outputs, 1)for j in range(inputs.size()[0]):images_so_far += 1ax = plt.subplot(num_images//2, 2, images_so_far)ax.axis('off')ax.set_title(f'predicted: {class_names[preds[j]]}')imshow(inputs.cpu().data[j])if images_so_far == num_images:model.train(mode=was_training)returnmodel.train(mode=was_training)

四、卷积网络微调

上三节都是准备工作，这一节，我们就讲一下两种迁移学习的使用。

微调 ConvNet：

#加载一个预训练的模型并且重置全连接层
model_ft = models.resnet18(weights='IMAGENET1K_V1')
num_ftrs = model_ft.fc.in_features
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to ``nn.Linear(num_ftrs, len(class_names))``.
model_ft.fc = nn.Linear(num_ftrs, 2)
model_ft = model_ft.to(device)
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_ft, criterion, optimizer_ft, exp_lr_scheduler,num_epochs=25)
visualize_model(model_ft)

ConvNet 作为固定特征提取器：

注意：这里，我们需要冻结除最后一层之外的所有网络。我们需要设置requires_grad = False来冻结参数，这样梯度就不会在backward()中计算。

model_conv = torchvision.models.resnet18(weights='IMAGENET1K_V1')
for param in model_conv.parameters():param.requires_grad = False# Parameters of newly constructed modules have requires_grad=True by default
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)model_conv = model_conv.to(device)
criterion = nn.CrossEntropyLoss()# Observe that only parameters of final layer are being optimized as
# opposed to before.
optimizer_conv = optim.SGD(model_conv.fc.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_conv, step_size=7, gamma=0.1)model_conv = train_model(model_conv, criterion, optimizer_conv, exp_lr_scheduler, num_epochs=25)
visualize_model(model_conv)
plt.ioff()
plt.show()

五、自定义图像的推理

使用训练的模型进行自定义图片的预测并且显示预测图片对应的标签

def visualize_model_predictions(model,img_path):was_training = model.trainingmodel.eval()img = Image.open(img_path)img = data_transforms['val'](img)img = img.unsqueeze(0)img = img.to(device)with torch.no_grad():outputs = model(img)_, preds = torch.max(outputs, 1)ax = plt.subplot(2,2,1)ax.axis('off')ax.set_title(f'Predicted: {class_names[preds[0]]}')imshow(img.cpu().data[0])model.train(mode=was_training)visualize_model_predictions(model_conv,img_path='data/hymenoptera_data/val/bees/72100438_73de9f17af.jpg'
)plt.ioff()
plt.show()