在Rust中使用TensorFlow

TensorFlow是一个开源的机器学习框架，上手简单，且在大量商业场景使用，可靠性高。

机器学习常用的语言主要是python、c，官方也提供了Java版本sdk，示例也很丰富。对于Rust语言，虽然官方文档没有写，其实官方也提供了rust版本的绑定，项目名称就是rust。

在Cargo.toml中配置版本，这里的0.18.0对应TensorFlow 2.8.0

tensorflow = "0.18.0"

这个依赖中的-sys模块会根据实际情况进行处理，如果本地没有Tensorflow的库文件，就会自动下载。如果没有预编译好的库文件，或者用户指定要从源码编译，那就会在本地编译。

我这里使用的是frozen graph，在rust代码中新建一个Graph进行加载即可。

let mut graph = Graph::new();
        let model_file = MODEL_DIR
            .get_file("mobilenet/mobilenet_v2_1.4_224_frozen.pb")
            .unwrap();
        let label_file = MODEL_DIR.get_file("mobilenet/label.txt").unwrap();
        graph
            .import_graph_def(model_file.contents(), &ImportGraphDefOptions::new())
            .unwrap();

为了最后打包方便，这里使用了include_dir包，会把模型和最终产物打包到一起。

这里的模型是一个mobilenet的图片分类模型，输入为224像素。图片的处理使用image包。

let img = image::open(photo.get_store_path())?;
        let resized = image::imageops::thumbnail(&img, 224, 224);
        let mut flattened: Vec<f32> = Vec::new();
        for rgb in resized.pixels() {
            flattened.push(rgb[0] as f32 / 255.);
            flattened.push(rgb[1] as f32 / 255.);
            flattened.push(rgb[2] as f32 / 255.);
        }
        {
            let input = Tensor::new(&[1, 224, 224, 3])
                .with_values(&flattened)
                .unwrap();
            let session = Session::new(&SessionOptions::new(), &graph)?;
            let mut args = SessionRunArgs::new();
            args.add_feed(
                &graph.operation_by_name_required("input").unwrap(),
                0,
                &input,
            );
            let prediction = args.request_fetch(
                &graph.operation_by_name_required("MobilenetV2/Predictions/Softmax")?,
                0,
            );
            session.run(&mut args).unwrap();
            let prediction_res: Tensor<f32> = args.fetch(prediction)?;
            let mut i = 0;
            let mut json_vec: Vec<f32> = Vec::new();
            while i < prediction_res.len() {
                json_vec.push(prediction_res[i]);
                i += 1;
            }
            let label = labels.get(imax(&json_vec).unwrap()).unwrap();
            return Ok(label.to_string());
        }

整体使用还是很简单的，只是部分文档缺失，需要研究下。如果有其他语言的Tensorflow使用经验会好很多。

目前Tensorflow Hub做的很好了，模型可以直接从上面下载，Savemodel也可以直接加载，也可以转为frozen后使用。

本文版权归作者所有，欢迎转载，但未经作者同意必须保留此段声明，且在文章页面明显位置给出原文连接，否则保留追究法律责任的权利。

转载自夜明的孤行灯

本文链接地址:https://www.huangyunkun.com/2022/08/15/use-tensorflow-with-rust/