Slow libtorch/tch-rs - Output Archive trace?

Kaiyotech: <p>Hello,</p>
<p>I’m working on moving some Python code into Rust, using tch-rs. I’m getting better performance than going back to Python using pyO3 bindings into my Rust, but only barely, which seems wrong.</p>
<p>When I do a trace, my thread is spending almost all of its time doing <code>torch::serialize::OutputArchive::save_to</code> and I’m not sure why, since I’m not intending to output  or save anything. I’m going to attempt to attach a call stack example.</p>
<p>Struct that includes the network and then the important bit of the new and then the bit with the forward pass.</p>
<pre><code class="lang-auto">pub struct GymWrapper { 
     gym: Gym, 
     net: Box&lt;dyn nn::Module&gt;, 
 }

impl GymWrapper{
pb fn new{
// other code to set up the gym removed...
        let var_store = nn::VarStore::new(Device::Cpu); 
        let net = network::net(&amp;var_store.root()); 
        GymWrapper { gym , net: Box::new(net)} 
     }

    pub fn step_episode(&amp;mut self, seed: Option&lt;u64&gt;) -&gt; bool { 
  
         tch::set_num_threads(1); 
  
         let mut steps = 0; 
         let mut done = false; 
         for _i in 0..10{ 
             done = false; 
             let mut obs = self.gym.reset(Some(false), seed); 
             
             while !done{ 
                 let tens_obs = Tensor::from_slice2(&amp;obs); 
                 let actions: Tensor = tch::no_grad(|| self.net.forward(&amp;tens_obs)); 
                 let act_vec: Vec&lt;Vec&lt;f32&gt;&gt; = Tensor::try_into(actions).expect("error from tensor to vector"); 
                 let result = self.gym.step(act_vec); 
                 obs = result.0;
                 done = result.2; 
                 steps += 1; 
             } 
  
         }
</code></pre>
<p>The test is to set up the gym wrapper and then time running the step episode. While step episode is running I took the attached trace. Like 30% of the exclusive time is in the <code>torch::serialize::OutputArchive::save_to</code>.</p>
<p>Any ideas what I’m doing wrong or what it’s trying to do?</p>
<p>Network is below because it’s long and ugly.</p>
<pre><code class="lang-auto">const HIDDEN_NODES: i64 = 256; 
 const INPUT_DIM: i64 = 231; 
 const OUTPUT_DIM: i64 = 90;


pub fn net(my_net: &amp;nn::Path) -&gt; impl Module { 
     nn::seq().add(nn::linear(my_net / "layer1", 
                 INPUT_DIM, 
                 HIDDEN_NODES, 
                 nn::LinearConfig{ws_init:Kaiming{dist: Normal, 
                     fan: FanIn, 
                    non_linearity: ReLU},bs_init:Some(Kaiming{dist: Normal, 
                      fan: FanIn, 
                     non_linearity: ReLU}),bias:true})) 
             .add_fn(|xs| xs.leaky_relu()) 
             .add(nn::linear(my_net / "layer2", 
                 HIDDEN_NODES, 
                 HIDDEN_NODES, 
                 nn::LinearConfig{ws_init:Kaiming{dist: Normal, 
                     fan: FanIn, 
                    non_linearity: ReLU},bs_init:Some(Kaiming{dist: Normal, 
                      fan: FanIn, 
                     non_linearity: ReLU}),bias:true})) 
             .add_fn(|xs| xs.leaky_relu()) 
             .add(nn::linear(my_net / "layer3", 
                 HIDDEN_NODES, 
                 HIDDEN_NODES, 
                 nn::LinearConfig{ws_init:Kaiming{dist: Normal, 
                     fan: FanIn, 
                    non_linearity: ReLU},bs_init:Some(Kaiming{dist: Normal, 
                      fan: FanIn, 
                     non_linearity: ReLU}),bias:true})) 
             .add_fn(|xs| xs.leaky_relu()) 
             .add(nn::linear(my_net / "layer4", 
                 HIDDEN_NODES, 
                 OUTPUT_DIM, 
                 nn::LinearConfig{ws_init:Kaiming{dist: Normal, 
                     fan: FanIn, 
                    non_linearity: ReLU},bs_init:Some(Kaiming{dist: Normal, 
                      fan: FanIn, 
                     non_linearity: ReLU}),bias:true})) 
 }
</code></pre>
<p>Trace here</p>
<aside class="onebox googledrive" data-onebox-src="https://drive.google.com/file/d/1XTtqq0rKyFAglWUUbOXlFLfff5JExOAN/view?usp=drivesdk">
  <header class="source">

      <a href="https://drive.google.com/file/d/1XTtqq0rKyFAglWUUbOXlFLfff5JExOAN/view?usp=drivesdk" target="_blank" rel="noopener nofollow ugc">drive.google.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://drive.google.com/file/d/1XTtqq0rKyFAglWUUbOXlFLfff5JExOAN/view?usp=drivesdk" target="_blank" rel="noopener nofollow ugc"><span class="googledocs-onebox-logo g-drive-logo"></span></a>



<h3><a href="https://drive.google.com/file/d/1XTtqq0rKyFAglWUUbOXlFLfff5JExOAN/view?usp=drivesdk" target="_blank" rel="noopener nofollow ugc">callstack_outputarchive_20230907.csv</a></h3>

<p>Google Drive file.</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>