### 加载自定义训练模型并执行推理 为了在 PyTorch 中加载自定义训练的模型并进行推理，可以遵循一系列特定的操作流程。这通常涉及几个关键步骤：导入必要的库、定义模型架构、加载保存的状态字典以及设置评估模式。 #### 导入库和工具 首先需要确保已经安装了 `torch` 和其他可能需要用到的相关包： ```bash pip install torch torchvision torchaudio ``` 接着，在脚本中引入这些模块： ```python import torch from PIL import Image import torchvision.transforms as transforms ``` #### 定义模型结构 假设有一个名为 `MyCustomModel` 的类来表示要加载的神经网络，则需先定义该类。如果是在同一项目内开发此模型的话可以直接从源文件导入；如果是外部资源则应按照其文档说明操作。这里给出一个简单的例子作为示意[^1]： ```python class MyCustomModel(torch.nn.Module): def __init__(self): super(MyCustomModel, self).__init__() # Define layers here... def forward(self, x): # Implement forward pass logic ... return output ``` #### 载入已训练权重 一旦有了模型实例之后就可以通过调用 `.load_state_dict()` 方法读取之前保存下来的参数值，并将其应用到当前创建的新对象上。注意路径应当指向实际存放`.pth` 或者 `.pt` 文件的位置[^2]: ```python model = MyCustomModel() model.load_state_dict(torch.load('path/to/custom_model_weights.pth')) model.eval() # Set model to evaluation mode. ``` 对于较大的预训练模型下载链接可参照给定命令行指令获取指定版本的数据集： ```bash mkdir -p models/ldm/text2img-large/ wget -O models/ldm/text2img-large/model.ckpt https://ommer-lab.com/files/latent-diffusion/nitro/txt2img-f8-large/model.ckpt ``` #### 执行预测任务 完成上述准备工作以后便能够利用输入数据来进行前向传播计算从而得到预期的结果。下面展示了一个完整的端到端过程，包括图片预处理在内的代码片段： ```python def preprocess_image(image_path): transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) image = Image.open(image_path) tensor = transform(image).unsqueeze(0) # Add batch dimension return tensor image_tensor = preprocess_image('example.jpg') with torch.no_grad(): # Disable gradient calculation during inference outputs = model(image_tensor) predicted_class_index = int(outputs.argmax()) print(f'Predicted class index: {predicted_class_index}') ```