“站在巨人肩膀上”——这正是深度学习中预训练模型所带来的好处。
从零开始训练深度神经网络既耗费资源又消耗时间，不过好消息是 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 是非常重要的，尤其是在使用预训练模型时，因为通常会处理高维数据和深层架构。
请确保 CUDA 和 cuDNN 已正确安装并可供 PyTorch 访问。以下是可以运行的快速检查命令：
import torch
# Check if GPU is available
print("CUDA available:", torch.cuda.is_available())
# Check CUDA version
print("CUDA version:", torch.version.cuda)
# Number of GPUs
print("Number of GPUs:", torch.cuda.device_count())对于多 GPU 配置，PyTorch 提供了 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel 来处理模型并行。以下是使用 DataParallel 的快速设置方法：
import torch
import torchvision.models as models
# Load model and enable DataParallel if multiple GPUs are available
model = 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 model
model = models.resnet50(pretrained=True)
model.eval() # Set to evaluation mode提示：在推理之前，请使用 model.eval()将模型设置为评估模式。
这可以确保像批归一化和丢弃层这样的层表现一致。
你可能想知道：“有哪些预训练模型可以直接使用？”PyTorch 提供了一种便捷的方法来直接从 torchvision 查看它们。
以下是列出所有可用模型的快速函数：
from torchvision import models
# List all available pretrained models in torchvision
available_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 architecture
model.load_state_dict(torch.load('custom_model_weights.pth')) # Load custom weights
model.eval() # Set to evaluation mode使用这种方法，可以为特定任务定制预训练模型，从而更容易地利用以前的工作，同时适应独特的需求。
在使用预训练模型时，微调是一种强大的工具，可以将模型学习到的特征适应到特定任务中。
以下是每种微调技术的结构化方法和实用代码片段。
想象一下：我们希望模型学习新特征，同时不“忘记”它已经掌握的基础特征。
为实现这一点，可以冻结特定层（特别是捕获通用特征的低层），同时保持高层可训练以进行微调。以下是操作方法：
import torch
import torchvision.models as models
# Load pretrained ResNet model
model = models.resnet50(pretrained=True)
# Freeze all layers by default
for param in model.parameters():
param.requires_grad = False
# Unfreeze the final fully connected layer
for 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 classes
num_classes = 10 # Adjust according to your dataset
model.fc = nn.Linear(model.fc.in_features, num_classes)你可能想知道：“为什么不修改模型的其他部分？”当然可以！但是，仅替换最后一层通常可以在性能和计算效率之间取得平衡。
在微调时，对不同部分的模型应用不同的学习率通常是有益的。
可以将较低的学习率分配给冻结或半冻结层，而将较高的学习率分配给新的分类器层。以下是如何在优化器中配置参数组：
# Define parameter groups with different learning rates
optimizer = 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 validation
train_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 torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm # For a progress bar
# Assuming `model` is defined, with `train_loader` and `val_loader` already prepared
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
# Early stopping parameters
early_stop_patience = 5
best_val_loss = float('inf')
patience_counter = 0
# Training loop
num_epochs = 20 # Set as needed
for 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 监控模型性能，同时使用提前停止来防止过拟合。
你将根据验证损失保存最佳模型，然后可以重新加载它进行进一步分析或部署。
让预训练模型运行得更快、更高效可能是改变游戏规则的关键，尤其是当你处理大型数据集或实时应用程序时。
以下是可以利用的一些强大优化技术。
混合精度训练结合了 16 位和 32 位浮点数来加速计算并节省内存。
PyTorch 的自动混合精度（AMP）功能可以无缝实现这一点。使用 AMP，你可以在保持模型准确性的同时减少训练时间和 GPU 内存使用量。以下是实现方法：
import torch
from torch.cuda.amp import GradScaler, autocast
# Assuming `model`, `optimizer`, `train_loader`, and `criterion` are already defined
scaler = 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 往往不够用，PyTorch 使用 torch.nn.DataParallel 使多 GPU 训练相对简单。
对于更大的配置，如TPU或分布式训练，torch.nn.parallel.DistributedDataParallel（DDP）是理想选择。
以下是设置多 GPU 的 DataParallel 快速示例：
# Assuming `model` is defined
model = torch.nn.DataParallel(model)
model = model.to(device)对于使用 DDP 的分布式训练，以下是基本设置：
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group("nccl") # Initialize the process group
model = DDP(model.to(device)) # Wrap model in DDP专业提示：使用 DDP 时，每个进程控制一个不同的 GPU，这可以显著减少大型数据集的训练时间。
如何最充分利用 GPU 内存？”答案在于选择合适的批量大小。
更大的批量大小可以加快训练速度，但存在内存过载的风险，尤其是对于高分辨率数据。
这里有个小技巧：从保守的批量大小开始，并逐步增加，同时监控 GPU 内存使用情况。
# Start with a small batch size and gradually increase
batch_size = 32
while 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_dict
torch.save(model.state_dict(), "model_weights.pth")
# Serialize with torch.jit for optimized inference
example_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（Open Neural Network Exchange）允许您在不同平台上部署模型。以下是将 PyTorch 模型导出为 ONNX 格式的快速流程：
import torch.onnx
# Define example input and export the model
example_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 等框架的兼容性，扩大了在不同硬件和环境中的部署选项。
实时推理实用技巧
如果要在实时应用程序中部署模型，需要尽可能减少延迟。
以下是一些有效技术：
1）批量推理：同时处理多个输入，这通常比单推理模式更快。
2）量化：使用 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 以及使用量化进行实时推理。
转自：