- 加载模型:使用
torch.load()
函数加载已经训练好的模型。 - 加载数据:使用
torchvision.datasets.ImageFolder()
函数加载待预测的imageFolder数据集。 - 预处理数据:使用
torchvision.transforms.Compose()
函数将待预测的数据集进行预处理,如将图片缩放、裁剪、标准化等。 - 预测数据:使用模型对待预测的数据进行预测,可以使用
model.eval()
函数将模型设置为评估模式,使用model.forward()
函数进行前向传播计算,得到预测结果。 - 输出结果:将预测结果输出,可以使用
print()
函数输出,也可以将结果保存到文件中。 下面是一个示例代码,假设已经训练好了一个分类模型model.pth
,待预测的数据集位于data/test/
目录下:
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)