tensorrt-rs 0.3.0

Rust library for using Nvidia's TensorRT deep learning acceleration library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

use crate::engine::Engine;
use crate::network::Network;
use crate::runtime::Logger;
use tensorrt_sys::{build_cuda_engine, create_infer_builder, create_network, destroy_builder};

pub struct Builder {
    internal_builder: *mut tensorrt_sys::Builder_t,
    network: Network,
}

impl Builder {
    pub fn new(logger: &Logger) -> Builder {
        let builder = unsafe { create_infer_builder(logger.internal_logger) };
        let sys_network = unsafe { create_network(builder) };
        let network = Network {
            internal_network: sys_network,
        };

        Builder {
            internal_builder: builder,
            network,
        }
    }

    pub fn get_network(&self) -> &Network {
        &self.network
    }

    pub fn build_cuda_engine(&self) -> Engine {
        Engine {
            internal_engine: unsafe {
                build_cuda_engine(self.internal_builder, self.network.internal_network)
            },
        }
    }
}

impl Drop for Builder {
    fn drop(&mut self) {
        unsafe { destroy_builder(self.internal_builder) };
    }
}