FastAI模型在训练完成之后,需要投入到生成环境中,用于新图片的预测。
1. 模型的保存和加载
模型的导出可以使用:
learn.export()
这个会在目录中创建一个名为export.pkl的文件,它包含了部署模型所需要的所有内容(模型,权重以及一些元数据)等,但我本人不太愿意用这种方法,因为我无法准确指定保存的路径和名称,这时,可以用:
learn.export('/home/ray/DataSet/models/model_resnet50_V190918.pkl')
模型的加载也可以使用:
learn = load_learner(path)
,此时会加载path中export.pkl这个模型,但是如果是我们自己定义了模型的保存路径,可以使用:
saved_learn=load_learner('/home/ray/DataSet/models',file='model_resnet50_V190918.pkl')
2. 模型的预测
对于新图片的预测,可以使用:
img = open_image(path/'grizzly bear/2.grizzly_pam-hartman_0_epv0457.jpg')
pred_class, pred_idx, outputs = learn.predict(img)
pred_class
所得到的pred_class就是预测出来的类别,outputs是预测出来的各个不同类别的概率。
在生产环境中,很多时候并没有GPU,所以我们要在CPU上进行预测,所以在预测之前,需要事先设置默认的设备:
defaults.device = torch.device('cpu')
网友评论