Pytorch官方提供的C++API名為libtorch蚕甥,在Pytorch1.0版本之后開始支持windows,以后就可以直接用libtorch來部署Pytorch模型了栋荸。
在Windows平臺上用C++部署深度學習應用一直很麻煩菇怀,caffe使用自定義層還需要自己寫backward(),而TensorFlow對Windows平臺的支持也不是很好晌块,在Windows上編譯有各種坑等著你爱沟。而這次Facebook提供的libtorch還是編譯好的文件,減少了很多的工作量匆背,估計會吸引一部分人轉向pytorch呼伸。
下面是我將pytorch的模型轉化為libtorch可用的torchscript模型的代碼。
import torch
import torchvision
# 加載模型
from models import FaceNetModel
embedding_size = 128
num_classes = 10000
model = FaceNetModel(embedding_size = embedding_size,num_classes = num_classes)
checkpoint = torch.load(r"D:\scripts\simple_facenet\log\checkpoint_epoch1610.pth")
model.load_state_dict(checkpoint['state_dict'])
model.cuda()
model.eval()
# 向模型中輸入數(shù)據(jù)以得到模型參數(shù)
example = torch.rand(1,3,182,182).cuda()
traced_script_module = torch.jit.trace(model,example)
# 保存模型
traced_script_module.save("torch_script_eval.pt")
這里有兩個地方容易出錯
- 如果模型中有dropout或者batchnorm的話钝尸,一定要先將模型設置為eval模式括享,再保存,否則在用libtorch調用后會出現(xiàn)隨機干擾珍促;
- example這個張量的尺寸務必與你自己的模型的輸入尺寸一直铃辖,否則會出現(xiàn)錯誤。
這個libtorch出來時間不長猪叙,還有很多坑等著我去踩娇斩,如果有同樣在研究pytorch或者libtorch的朋友,可以加個好友沐悦,共同學習成洗。