使用训练好的PyTorch模型进行预测可以分为以下步骤:
- 加载模型:使用
torch.load()
函数加载已经训练好的模型。 - 加载数据:使用
torchvision.datasets.ImageFolder()
函数加载待预测的imageFolder数据集。 - 预处理数据:使用
torchvision.transforms.Compose()
函数将待预测的数据集进行预处理,如将图片缩放、裁剪、标准化等。 - 预测数据:使用模型对待预测的数据进行预测,可以使用
model.eval()
函数将模型设置为评估模式,使用model.forward()
函数进行前向传播计算,得到预测结果。 - 输出结果:将预测结果输出,可以使用
print()
函数输出,也可以将结果保存到文件中。 下面是一个示例代码,假设已经训练好了一个分类模型model.pth
,待预测的数据集位于data/test/
目录下:
pythonCopy code
import torch import torchvision.transforms as transforms import torchvision.datasets as datasets # 加载模型 model = torch.load('model.pth') # 加载数据 data_transforms = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_data = datasets.ImageFolder(root='data/test', transform=data_transforms) # 预测数据 model.eval() with torch.no_grad(): for inputs, _ in test_data: inputs = inputs.unsqueeze(0) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) print(predicted)
在以上代码中,首先使用torch.load()
函数加载已经训练好的模型,然后使用torchvision.datasets.ImageFolder()
函数加载待预测的数据集。接着,使用torchvision.transforms.Compose()
函数将待预测的数据集进行预处理,包括将图片缩放、裁剪、标准化等操作。然后,将模型设置为评估模式,使用model.forward()
函数进行前向传播计算,得到预测结果。最后,将预测结果输出。