放心加上cuda 0 也没关系,代码中自动检测是否有显卡,没有自动换成cpu.
保存为onnx后你需要转换成coreML
去年一个师兄毕业,在手机上做了一个app,分图像中的狗和猫。早上没什么事,就想着在ios上玩玩风格迁移。
由于疫情原因,手里没机器就去下载了一个,我用的是这个https://github.com/pytorch/examples/tree/master/fast_neural_style
他的模型里面保存的是参数,而且是pth,ios上得用coreML,所以得用Open Neural Network Exchange(ONNX)转换一下。
当然也不用你自己转换,上面的包中有onnx
python ./neural_style/neural_style.py eval \
--content-image ./images/original.jpg \
--output-image ./stylized-test.jpg \
--model ./saved_model/candy.pth \
--cuda 0 \
--export_onnx ./candy.onnx
然后将onnx转换成coreML
# onnx2coreML.py
import sys
from onnx import onnx_pb
from onnx_coreml import convert
model_in = sys.argv[1]
model_out = sys.argv[2]
model_file = open(model_in, 'rb')
model_proto = onnx_pb.ModelProto()
model_proto.ParseFromString(model_file.read())
coreml_model = convert(model_proto, image_input_names=['inputImage'], image_output_names=['outputImage'])
coreml_model.save(model_out)
python onnx2coreML.py candy.onnx candy.mlmodel
到这里就有了可以在ios上运行的模型了。
无奈不会swift所以上网找了一个https://github.com/onnx/tutorials
发现https://github.com/onnx/tutorials/tree/master/examples/CoreML/ONNXLive里面什么都有。。。。。直接往ios部署就好。。。。
总的来说,onnx里面都有,只要你部署下就能出去装X了,醉醉的。。。。。