From 089aab793046d6dce29a6568056d4fcb1ecdb60d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 13 Oct 2019 21:23:11 -0400 Subject: [PATCH] move model stuff to new file --- src/main.rs | 168 +++++---------------------------------------------- src/model.rs | 141 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 154 deletions(-) create mode 100644 src/model.rs diff --git a/src/main.rs b/src/main.rs index 75a3483208..db857cb3d9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,8 +6,10 @@ extern crate tensorflow; extern crate wavy; extern crate world_sys; extern crate zmq; +extern crate failure; pub mod ipc; +pub mod model; pub mod synthesis; use wavy::*; @@ -19,22 +21,13 @@ use std::path::Path; use std::process::exit; use std::thread; -use tensorflow::Code; -use tensorflow::Graph; -use tensorflow::ImportGraphDefOptions; -use tensorflow::Session; -use tensorflow::SessionOptions; -use tensorflow::SessionRunArgs; -use tensorflow::Status; -use tensorflow::Tensor; -use tensorflow::version; use zmq::{Error, Socket}; use ipc::{QueueSender, AudioQueue}; +use model::load_model; +use model::print_version; use std::sync::Arc; - -const INPUT_NAME : &'static str = "input_A_test"; -const OUTPUT_NAME : &'static str = "generator_A2B_3/output_transpose"; +use failure::_core::time::Duration; fn main() { print_version(); @@ -43,135 +36,6 @@ fn main() { //run_audio().expect("should work"); } -fn print_version() { - // Python TensorFlow version: 1.14.0 - // Rust TensorFlow version: 1.13.1 (default, must be hand-upgraded) - let version = version().expect("version"); - println!("Tensorflow version: {}", version); -} - -struct VoiceModel { - graph: Graph, - session: Session, -} - -impl VoiceModel { - pub fn load(dir_name: &str) -> Self { - let mut graph = Graph::new(); - let session = Session::from_saved_model( - &SessionOptions::new(), - &["serve"], - &mut graph, - dir_name, - ).expect("Should load"); - - VoiceModel { - graph, - session, - } - } - - pub fn evaluate(&self, input: &Tensor) { - println!(">>> Input tensor dims: {:?}", input.dims()); - - // Run the graph. - let mut args = SessionRunArgs::new(); - - // input_A_test: - // Tensor("input_A_test:0", shape=(?, 24, ?), dtype=float32) - args.add_feed(&self.graph.operation_by_name_required(INPUT_NAME) - .expect(INPUT_NAME), 0, &input); - - // generation_B_test: - // Tensor("generator_A2B_3/output_transpose:0", shape=(?, 24, ?), dtype=float32) - let z = args.request_fetch( - &self.graph.operation_by_name_required(OUTPUT_NAME) - .expect(OUTPUT_NAME), 0); - - self.session.run(&mut args).expect("Run success"); - - // Check our results. - let z_res = args.fetch::(z).expect("ret"); - - println!("z_rez.dims(): {:?}", z_res.dims()); - println!("z_rez: {:?}", z_res); - } -} - - -fn load_model() { - let model_dir = "/home/bt/dev/voder/saved_model"; - let model = VoiceModel::load(model_dir); - - // ] generation_B_test: Tensor("generator_A2B_3/output_transpose:0", shape=(?, 24, ?), dtype=float32) - // ] input_A_test: Tensor("input_A_test:0", shape=(?, 24, ?), dtype=float32) - // ] inputs.shape: (1, 24, 380) -- the last dimension is the signal - - let mut input = Tensor::new(&[1, 24, 10]) - .with_values(&[ - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - ]) - .unwrap(); - - model.evaluate(&input); - model.evaluate(&input); - model.evaluate(&input); -} - -/*fn load_model_2() { - // from regression_checkpoint.rs example - - let filename = "/home/bt/dev/voder/saved_model/saved_model.pb"; - - let mut graph = Graph::new(); - let mut proto = Vec::new(); - - File::open(filename) - .expect("opened") - .read_to_end(&mut proto) - .expect("cannot read"); - - graph.import_graph_def(&proto, &ImportGraphDefOptions::new()).expect("cannot import"); - - let session = Session::new(&SessionOptions::new(), &graph).expect("cannot session"); - - /*let op_x = graph.operation_by_name_required("x").expect("x"); - let op_y = graph.operation_by_name_required("y").expect("y"); - let op_init = graph.operation_by_name_required("init").expect("init"); - let op_train = graph.operation_by_name_required("train").expect("train"); - let op_w = graph.operation_by_name_required("w").expect("w"); - let op_b = graph.operation_by_name_required("b").expect("b"); - let op_file_path = graph.operation_by_name_required("save/Const").expect("const"); - let op_save = graph.operation_by_name_required("save/control_dependency").expect("save");*/ - - //let file_path_tensor: Tensor = - // Tensor::from(String::from("/home/bt/dev/voder/extra_model/sf1_tm1.ckpt.data-00000-of-00001")); - //println!("Tensor: {:?}", file_path_tensor); -}*/ - fn run_audio() -> Result<(), AudioError> { println!("Opening microphone system"); let mut mic = MicrophoneSystem::new(SampleRate::Normal)?; @@ -187,19 +51,13 @@ fn run_audio() -> Result<(), AudioError> { let mut audio_queue_2 = audio_queue.clone(); thread::spawn(move || { - //let mut queue_sender = QueueSender::new(); - //queue_sender.connect(); - let mut context = zmq::Context::new(); let mut socket = context.socket(zmq::REQ).unwrap(); - /*let mut socket = ctx.socket(zmq::REQ).unwrap(); - socket.connect("tcp://127.0.0.1:1234").unwrap(); - socket.send("hello world!", 0).unwrap();*/ - socket.connect("tcp://127.0.0.1:5555").unwrap(); let mut reconnect = false; + let mut fail_count = 0; loop { let mut drained = audio_queue_2.drain(); @@ -211,6 +69,8 @@ fn run_audio() -> Result<(), AudioError> { if reconnect { reconnect = false; + thread::sleep(Duration::from_millis(200)); + socket = match context.socket(zmq::REQ) { Ok(s) => s, Err(e) => { @@ -220,23 +80,23 @@ fn run_audio() -> Result<(), AudioError> { }; match socket.connect("tcp://127.0.0.1:5555") { - Ok(_) => { - println!("B"); - }, + Ok(_) => {}, Err(err) => { println!("Err B: {:?}", err); }, } - } match socket.send(&bytes[..], 0) { Ok(_) => { break; }, Err(e) => { println!("send err: {:?}", e); - reconnect = true; + fail_count += 1; + if fail_count > 5 { + fail_count = 0; + reconnect = true; + } }, } - } } } diff --git a/src/model.rs b/src/model.rs new file mode 100644 index 0000000000..b9a982fa27 --- /dev/null +++ b/src/model.rs @@ -0,0 +1,141 @@ +use tensorflow::Code; +use tensorflow::Graph; +use tensorflow::ImportGraphDefOptions; +use tensorflow::Session; +use tensorflow::SessionOptions; +use tensorflow::SessionRunArgs; +use tensorflow::Status; +use tensorflow::Tensor; +use tensorflow::version; + +const INPUT_NAME : &'static str = "input_A_test"; +const OUTPUT_NAME : &'static str = "generator_A2B_3/output_transpose"; + +pub struct VoiceModel { + graph: Graph, + session: Session, +} + +impl VoiceModel { + pub fn load(dir_name: &str) -> Self { + let mut graph = Graph::new(); + let session = Session::from_saved_model( + &SessionOptions::new(), + &["serve"], + &mut graph, + dir_name, + ).expect("Should load"); + + VoiceModel { + graph, + session, + } + } + + pub fn evaluate(&self, input: &Tensor) { + println!(">>> Input tensor dims: {:?}", input.dims()); + + // Run the graph. + let mut args = SessionRunArgs::new(); + + // input_A_test: + // Tensor("input_A_test:0", shape=(?, 24, ?), dtype=float32) + args.add_feed(&self.graph.operation_by_name_required(INPUT_NAME) + .expect(INPUT_NAME), 0, &input); + + // generation_B_test: + // Tensor("generator_A2B_3/output_transpose:0", shape=(?, 24, ?), dtype=float32) + let z = args.request_fetch( + &self.graph.operation_by_name_required(OUTPUT_NAME) + .expect(OUTPUT_NAME), 0); + + self.session.run(&mut args).expect("Run success"); + + // Check our results. + let z_res = args.fetch::(z).expect("ret"); + + println!("z_rez.dims(): {:?}", z_res.dims()); + println!("z_rez: {:?}", z_res); + } +} + +pub fn load_model() { + let model_dir = "/home/bt/dev/voder/saved_model"; + let model = VoiceModel::load(model_dir); + + // ] generation_B_test: Tensor("generator_A2B_3/output_transpose:0", shape=(?, 24, ?), dtype=float32) + // ] input_A_test: Tensor("input_A_test:0", shape=(?, 24, ?), dtype=float32) + // ] inputs.shape: (1, 24, 380) -- the last dimension is the signal + + let mut input = Tensor::new(&[1, 24, 10]) + .with_values(&[ + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + ]) + .unwrap(); + + model.evaluate(&input); + model.evaluate(&input); + model.evaluate(&input); +} + +/*fn load_model_2() { + // from regression_checkpoint.rs example + + let filename = "/home/bt/dev/voder/saved_model/saved_model.pb"; + + let mut graph = Graph::new(); + let mut proto = Vec::new(); + + File::open(filename) + .expect("opened") + .read_to_end(&mut proto) + .expect("cannot read"); + + graph.import_graph_def(&proto, &ImportGraphDefOptions::new()).expect("cannot import"); + + let session = Session::new(&SessionOptions::new(), &graph).expect("cannot session"); + + /*let op_x = graph.operation_by_name_required("x").expect("x"); + let op_y = graph.operation_by_name_required("y").expect("y"); + let op_init = graph.operation_by_name_required("init").expect("init"); + let op_train = graph.operation_by_name_required("train").expect("train"); + let op_w = graph.operation_by_name_required("w").expect("w"); + let op_b = graph.operation_by_name_required("b").expect("b"); + let op_file_path = graph.operation_by_name_required("save/Const").expect("const"); + let op_save = graph.operation_by_name_required("save/control_dependency").expect("save");*/ + + //let file_path_tensor: Tensor = + // Tensor::from(String::from("/home/bt/dev/voder/extra_model/sf1_tm1.ckpt.data-00000-of-00001")); + //println!("Tensor: {:?}", file_path_tensor); +}*/ + +pub fn print_version() { + // Python TensorFlow version: 1.14.0 + // Rust TensorFlow version: 1.13.1 (default, must be hand-upgraded) + let version = version().expect("version"); + println!("Tensorflow version: {}", version); +} +