### 如何在 PyTorch 中加载并运行预训练的分类模型 #### 使用 Torchvision 加载预训练模型 PyTorch 的 `torchvision` 库提供了一系列常用的预训练模型，这些模型已经在大规模数据集上进行了训练。通过简单的几行代码就可以获取到这些预训练好的模型。 ```python import torch from torchvision import models # 获取 ResNet-50 预训练模型 model = models.resnet50(pretrained=True) # 切换至评估模式 model.eval() ``` 上述代码展示了如何从 `torchvision.models` 下载一个预先训练过的 ResNet-50 模型，并将其设置为评估模式以便于推理阶段使用[^1]。 #### 处理输入图像 为了让模型能够处理新图片，在将图片送入模型之前需要做一些准备工作： ```python from PIL import Image from torchvision import transforms input_image = Image.open('path_to_your_image.jpg') preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) input_tensor = preprocess(input_image) input_batch = input_tensor.unsqueeze(0) # 创建批次维度 (batch dimension) if torch.cuda.is_available(): input_batch = input_batch.to('cuda') model.to('cuda') with torch.no_grad(): # 关闭梯度计算以节省内存 output = model(input_batch) probabilities = torch.nn.functional.softmax(output[0], dim=0) print(probabilities.argmax().item()) # 输出预测类别索引 ``` 这段脚本说明了怎样准备一张新的测试图片以及执行前向传播得到最终的结果。注意这里关闭了自动求导功能 (`torch.no_grad`) 来减少不必要的运算开销[^3]。 #### 自定义网络结构与迁移学习 当自定义的神经网络架构不同于标准库里的时候，则可以采用迁移学习的方式利用已有的权重初始化部分层。具体做法如下所示: ```python class Net(torch.nn.Module): def __init__(self): super().__init__() self.backbone = ... # 定义骨干网或其他组件 def forward(self,x): ... # 初始化自定义模型实例 custom_model = Net() # 载入预训练模型字典 pretrained_dict = model.state_dict() # 过滤掉不匹配的部分 filtered_pretrained_dict = {k:v for k,v in pretrained_dict.items() if k in custom_model.state_dict()} # 更新自定义模型的状态字典 custom_model.load_state_dict(filtered_pretrained_dict, strict=False) ``` 此段代码片段解释了如果想要继承某些特定模块或层的情况下应该如何操作。通过这种方式可以在保持原有性能的同时调整适合目标任务的新设计[^2]。