学深度学习一定要知道的——pytorch 预训练模型

“站在巨人肩膀上”——这正是深度学习中预训练模型所带来的好处。

从零开始训练深度神经网络既耗费资源又消耗时间，不过好消息是 PyTorch 的预训练模型可以很好的解决这个问题，pytorch 可以提供已在大型数据集上学习到鲁棒特征的模型，让我们可以快速将这些模型适配到特定任务中，跳过繁重的训练工作，直接进行微调。

‍今天小墨就将带大家了解pytorch的预训练模型以及注意事项。

环境设置

本文默认大家已经设置好环境，将直接介绍让预训练模型在你的环境中顺畅运行的关键步骤。

‍

所需库

首先，请确保你已安装最新版本的 PyTorch 和 torchvision。以下是快速安装它们的代码片段：

# Install PyTorch and torchvision with GPU support# You can specify the appropriate CUDA version as needed.!pip install torch torchvision

‍

注意：为了获得最佳性能，请确保 PyTorch 的安装版本与 CUDA 版本相匹配。

机器上需要安装 CUDA 和 cuDNN，尤其是多 GPU 配置时。

GPU 设置

为深度学习任务正确配置 GPU 是非常重要的，尤其是在使用预训练模型时，因为通常会处理高维数据和深层架构。

请确保 CUDA 和 cuDNN 已正确安装并可供 PyTorch 访问。以下是可以运行的快速检查命令：

import torch
# Check if GPU is availableprint("CUDA available:", torch.cuda.is_available())# Check CUDA versionprint("CUDA version:", torch.version.cuda)# Number of GPUsprint("Number of GPUs:", torch.cuda.device_count())对于多 GPU 配置，PyTorch 提供了 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel 来处理模型并行。以下是使用 DataParallel 的快速设置方法：import torchimport torchvision.models as models
# Load model and enable DataParallel if multiple GPUs are availablemodel = models.resnet50(pretrained=True)if torch.cuda.device_count() > 1:    model = torch.nn.DataParallel(model)model = model.to('cuda')

完成此设置后，就可以充分利用硬件资源，并从预训练模型中获取最大收益。

加载预训练模型

现在，让我们进入正题：加载预训练模型。PyTorch 的 torchvision.models 库提供了在 ImageNet 上训练的多种模型，从 ResNet 和 VGG 到 MobileNet 和 EfficientNet。以下是高效加载它们的方法。

基本加载代码

首先，让我们从加载一个基本的预训练模型开始，ResNet50 是处理图像数据一个热门选择。以下是单行加载代码：

import torchvision.models as models
# Load a pretrained ResNet50 modelmodel = models.resnet50(pretrained=True)model.eval()  # Set to evaluation mode

提示：在推理之前，请使用 model.eval()将模型设置为评估模式。

这可以确保像批归一化和丢弃层这样的层表现一致。

可用模型列表

你可能想知道：“有哪些预训练模型可以直接使用？”PyTorch 提供了一种便捷的方法来直接从 torchvision 查看它们。

以下是列出所有可用模型的快速函数：

from torchvision import models
# List all available pretrained models in torchvisionavailable_models = [name for name in dir(models) if callable(getattr(models, name))]print("Available pretrained models:", available_models)

当你需要探索选项或在不同环境中工作时，此代码片段非常有用，它可以让你快速了解可用的模型。

自定义模型加载

有时，我们可能希望超越默认值，并加载具有自定义权重的模型。

如果在独特的数据集上进行了训练，或者正在使用 torchvision 中未直接提供的模型，这将特别有用。

以下是如何加载具有特定权重的自定义模型：

# Assuming you have a saved state dictionary 'custom_model_weights.pth'model = models.resnet50()  # Load model architecturemodel.load_state_dict(torch.load('custom_model_weights.pth'))  # Load custom weightsmodel.eval()  # Set to evaluation mode

使用这种方法，可以为特定任务定制预训练模型，从而更容易地利用以前的工作，同时适应独特的需求。

微调技术

在使用预训练模型时，微调是一种强大的工具，可以将模型学习到的特征适应到特定任务中。

以下是每种微调技术的结构化方法和实用代码片段。

冻结层

想象一下：我们希望模型学习新特征，同时不“忘记”它已经掌握的基础特征。

为实现这一点，可以冻结特定层（特别是捕获通用特征的低层），同时保持高层可训练以进行微调。以下是操作方法：

import torchimport torchvision.models as models
# Load pretrained ResNet modelmodel = models.resnet50(pretrained=True)
# Freeze all layers by default


    
for param in model.parameters():    param.requires_grad = False
# Unfreeze the final fully connected layerfor param in model.fc.parameters():    param.requires_grad = True

使用这种方法，我们只更新最后一层的参数，对模型进行微调，同时不干扰其更深层、已学习的表示。

调整模型架构

在许多实际应用中，预训练模型中的默认分类器头可能不适合您的任务。

例如，ResNet 的默认分类器是针对 ImageNet 的 1000 类设计的，但可能只需要几个类。以下是如何替换最后一层以使用自定义分类器：

import torch.nn as nn
# Modify the fully connected layer for a custom number of classesnum_classes = 10  # Adjust according to your datasetmodel.fc = nn.Linear(model.fc.in_features, num_classes)

你可能想知道：“为什么不修改模型的其他部分？”当然可以！但是，仅替换最后一层通常可以在性能和计算效率之间取得平衡。

优化器中的参数组

在微调时，对不同部分的模型应用不同的学习率通常是有益的。

可以将较低的学习率分配给冻结或半冻结层，而将较高的学习率分配给新的分类器层。以下是如何在优化器中配置参数组：

# Define parameter groups with different learning ratesoptimizer = torch.optim.SGD([    {'params': model.fc.parameters(), 'lr': 1e-3},  # Higher LR for new layers    {'params': model.layer4.parameters(), 'lr': 1e-4},  # Slightly lower LR    {'params': model.layer1.parameters(), 'lr': 1e-5}  # Lowest LR for earlier layers], momentum=0.9)

此设置确保模型受益于不同的学习率，帮助分类器更快地学习，同时保持早期层的稳定性。

实践应用：在自定义数据集上微调

将微调应用于与任务相关的数据集时，其意义更为重大。让我们了解如何准备自定义数据集并设置高效的训练循环。

数据集准备

为了提高效率，我们将使用 torchvision 的 ImageFolder 来加载自定义数据集。

以下代码片段展示了如何加载和应用基本变换以提高模型性能：

from torchvision import datasets, transforms
# Define transformations for training and validationtrain_transform = transforms.Compose([    transforms.Resize((224, 224)),    transforms.RandomHorizontalFlip(),    transforms.ToTensor(),    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
train_dataset = datasets.ImageFolder(root='path/to/train_data', transform=train_transform)train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

训练代码

我们将设置一个训练循环，其中包括：

训练和验证的损失和准确性记录。
提前停止以防止过拟合。
模型检查点以保存性能最佳的模型。

import torchimport torch.nn as nnimport torch.optim as optimfrom tqdm import tqdm  # For a progress bar
# Assuming `model` is defined, with `train_loader` and `val_loader` already prepareddevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = model.to(device)
# Loss function and optimizercriterion = nn.CrossEntropyLoss()optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
# Early stopping parametersearly_stop_patience = 5best_val_loss = float('inf')patience_counter = 0
# Training loopnum_epochs = 20  # Set as neededfor epoch in range(num_epochs):    model.train()  # Set model to training mode    running_loss = 0.0    correct = 0    total = 0
    # Training phase    for images, labels in tqdm(train_loader, desc=f"Training Epoch {epoch+1}"):        images, labels = images.to(device), labels.to(device)                optimizer.zero_grad()  # Zero the parameter gradients        outputs = model(images)        loss = criterion(outputs, labels)        loss.backward()        optimizer.step()
        running_loss += loss.item() * images.size(0)        _, predicted = outputs.max(1)        correct += predicted.eq(labels).sum().item()        total += labels.size(0)


    

    train_loss = running_loss / total    train_accuracy = 100. * correct / total    print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {train_loss:.4f}, Training Accuracy: {train_accuracy:.2f}%")
    # Validation phase    model.eval()    val_loss = 0.0    correct = 0    total = 0    with torch.no_grad():        for images, labels in val_loader:            images, labels = images.to(device), labels.to(device)            outputs = model(images)            loss = criterion(outputs, labels)
            val_loss += loss.item() * images.size(0)            _, predicted = outputs.max(1)            correct += predicted.eq(labels).sum().item()            total += labels.size(0)
    val_loss /= total    val_accuracy = 100. * correct / total    print(f"Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")
    # Early stopping and checkpointing    if val_loss < best_val_loss:        best_val_loss = val_loss        patience_counter = 0        torch.save(model.state_dict(), "best_model.pth")  # Save the best model        print("Best model saved!")    else:        patience_counter += 1        if patience_counter >= early_stop_patience:            print("Early stopping triggered. Ending training.")            break

此循环允许每个 epoch 监控模型性能，同时使用提前停止来防止过拟合。

你将根据验证损失保存最佳模型，然后可以重新加载它进行进一步分析或部署。

大家只要想结合机器学习深度学习技术做创新发论文，却没有头绪不知道怎么学习的，都可以直接添加助手微信号 “ai0808q ”了解详情（通过后回复咨询课程）

性能优化技术

让预训练模型运行得更快、更高效可能是改变游戏规则的关键，尤其是当你处理大型数据集或实时应用程序时。

以下是可以利用的一些强大优化技术。

混合精度训练

混合精度训练结合了 16 位和 32 位浮点数来加速计算并节省内存。

PyTorch 的自动混合精度（AMP）功能可以无缝实现这一点。使用 AMP，你可以在保持模型准确性的同时减少训练时间和 GPU 内存使用量。以下是实现方法：

import torchfrom torch.cuda.amp import GradScaler, autocast
# Assuming `model`, `optimizer`, `train_loader`, and `criterion` are already definedscaler = GradScaler()  # Initializes a gradient scaler for AMP
for epoch in range(num_epochs):    model.train()    for images, labels in train_loader:        images, labels = images.to(device), labels.to(device)              optimizer.zero_grad()              with autocast():  # Enables mixed precision            outputs = model(images)            loss = criterion(outputs, labels)              scaler.scale(loss).backward()  # Scales the loss to manage gradients safely        scaler.step(optimizer)        scaler.update()

为何有效：AMP（自动混合精度）利用自动类型转换功能为每个操作自动选择适当的精度，从而在保持模型准确性的同时节省内存。

多 GPU 或 TPU 配置

处理大型数据集时，单个 GPU 往往不够用，PyTorch 使用 torch.nn.DataParallel 使多 GPU 训练相对简单。

对于更大的配置，如TPU或分布式训练，torch.nn.parallel.DistributedDataParallel（DDP）是理想选择。

以下是设置多 GPU 的 DataParallel 快速示例：

# Assuming `model` is definedmodel = torch.nn.DataParallel(model)model = model.to(device)

对于使用 DDP 的分布式训练，以下是基本设置：

import torch.distributed as distfrom torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group("nccl")  # Initialize the process groupmodel = DDP(model.to(device))  # Wrap model in DDP

专业提示：使用 DDP 时，每个进程控制一个不同的 GPU，这可以显著减少大型数据集的训练时间。

批量大小缩放

如何最充分利用 GPU 内存？”答案在于选择合适的批量大小。

更大的批量大小可以加快训练速度，但存在内存过载的风险，尤其是对于高分辨率数据。

这里有个小技巧：从保守的批量大小开始，并逐步增加，同时监控 GPU 内存使用情况。

# Start with a small batch size and gradually increasebatch_size = 32while True:    try:


    
        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)        # Test one pass        for images, _ in train_loader:            images = images.to(device)        break    except RuntimeError:        batch_size //= 2  # Reduce batch size if memory overloads        print(f"Reduced batch size to {batch_size}")

为何有效：这种缩放技术可以最大化您的 GPU 内存，减少整体训练时间，同时避免内存崩溃的风险。

在生产中部署预训练模型

模型训练完成后，部署是下一步。以下是设置模型序列化、导出为 ONNX 以及配置实时推理的方法。

模型序列化

为了有效保存和加载模型，PyTorch 提供了 torch.save 来保存模型检查点。此外，torch.jit.trace 可用于优化模型，以加快推理速度：

# Save the model's state_dicttorch.save(model.state_dict(), "model_weights.pth")
# Serialize with torch.jit for optimized inferenceexample_input = torch.randn(1, 3, 224, 224).to(device)traced_model = torch.jit.trace(model, example_input)traced_model.save("traced_model.pt")

提示：torch.jit.trace 可以减少生产中的加载时间，使推理更快、更高效。

导出为 ONNX

ONNX（Open Neural Network Exchange）允许您在不同平台上部署模型。以下是将 PyTorch 模型导出为 ONNX 格式的快速流程：

import torch.onnx
# Define example input and export the modelexample_input = torch.randn(1, 3, 224, 224).to(device)torch.onnx.export(model, example_input, "model.onnx", export_params=True,                  input_names=["input"], output_names=["output"])

这实现了与 TensorFlow 和 ONNX Runtime 等框架的兼容性，扩大了在不同硬件和环境中的部署选项。

实时推理实用技巧

如果要在实时应用程序中部署模型，需要尽可能减少延迟。

以下是一些有效技术：

批量推理：同时处理多个输入，这通常比单推理模式更快。
量化：使用 torch.quantization 将模型权重转换为 int8，这可以显著减少内存并提高推理速度。

量化示例：

model.qconfig = torch.quantization.get_default_qconfig('fbgemm')model = torch.quantization.prepare(model)model = torch.quantization.convert(model)

通过这些技术，你的模型将针对生产进行优化，实现最低延迟。

结论

以下是本文的内容概览：

设置：从加载模型到配置多 GPU 环境。
微调：冻结层、调整架构和使用自定义参数组。
优化：混合精度训练、分布式设置和批量缩放以提高效率。
部署：使用 torch.jit 进行序列化、导出为 ONNX 以及使用量化进行实时推理。

‍

大家只要想结合机器学习深度学习技术做创新发论文，却没有头绪不知道怎么学习的，都可以直接添加助手微信号 “ai0808q ”了解详情（通过后回复咨询课程）

另外我还给大家准备了一些机器学习、深度学习、神经网络资料大家可以看看以下文章（文章中提到的资料都打包好了，都可以直接添加小助手获取）

人工智能资料分享

人工智能数学（点击图片即可跳转）

深度学习中文教程书（点击图片即可跳转）

神经网络最全学习资料（点击图片即可跳转）

大家觉得这篇文章有帮助的话记得分享给你的死党、闺蜜、同学、朋友、老师、敌蜜！

B站：AI秃秃学长小墨

关注小墨

获取最新AI技能+最肝AI干货