Reduce memory allocations (#350)

huggingface · Aug 30, 2023 · 5af8cba · 5af8cba
1 parent 73db0c8
commit 5af8cba
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
@@ -23,7 +23,7 @@ static TENSORFLOW_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 static FLAX_MODULE: GILOnceCell<Py<PyModule>> = GILOnceCell::new();
 
 fn prepare(tensor_dict: HashMap<String, &PyDict>) -> PyResult<HashMap<String, TensorView<'_>>> {
-    let mut tensors = HashMap::new();
+    let mut tensors = HashMap::with_capacity(tensor_dict.len());
     for (tensor_name, tensor_desc) in tensor_dict {
         let mut shape: Option<Vec<usize>> = None;
         let mut dtype: Option<Dtype> = None;
@@ -144,19 +144,21 @@ fn serialize_file(
 fn deserialize(py: Python, bytes: &[u8]) -> PyResult<Vec<(String, HashMap<String, PyObject>)>> {
     let safetensor = SafeTensors::deserialize(bytes)
         .map_err(|e| SafetensorError::new_err(format!("Error while deserializing: {e:?}")))?;
-    let mut items = vec![];
 
-    for (tensor_name, tensor) in safetensor.tensors() {
-        let mut map = HashMap::new();
+    let tensors = safetensor.tensors();
+    let mut items = Vec::with_capacity(tensors.len());
 
+    for (tensor_name, tensor) in tensors {
         let pyshape: PyObject = PyList::new(py, tensor.shape().iter()).into();
         let pydtype: PyObject = format!("{:?}", tensor.dtype()).into_py(py);
 
         let pydata: PyObject = PyByteArray::new(py, tensor.data()).into();
 
-        map.insert("shape".to_string(), pyshape);
-        map.insert("dtype".to_string(), pydtype);
-        map.insert("data".to_string(), pydata);
+        let map = HashMap::from([
+            ("shape".to_string(), pyshape),
+            ("dtype".to_string(), pydtype),
+            ("data".to_string(), pydata),
+        ]);
         items.push((tensor_name, map));
     }
     Ok(items)

diff --git a/safetensors/src/slice.rs b/safetensors/src/slice.rs
@@ -260,7 +260,8 @@ impl<'data> SliceIterator<'data> {
                         indices.push((offset, offset + small_span));
                     }
                 } else {
-                    let mut newindices = vec![];
+                    let capacity = (stop - start) * indices.len();
+                    let mut newindices = Vec::with_capacity(capacity);
                     for n in start..stop {
                         let offset = n * span;
                         for (old_start, old_stop) in &indices {

diff --git a/safetensors/src/tensor.rs b/safetensors/src/tensor.rs
@@ -349,7 +349,7 @@ impl<'data> SafeTensors<'data> {
     /// The tensors returned are merely views and the data is not owned by this
     /// structure.
     pub fn tensors(&self) -> Vec<(String, TensorView<'_>)> {
-        let mut tensors = vec![];
+        let mut tensors = Vec::with_capacity(self.metadata.index_map.len());
         for (name, &index) in &self.metadata.index_map {
             let info = &self.metadata.tensors[index];
             let tensorview = TensorView {
@@ -464,7 +464,7 @@ impl Metadata {
         metadata: Option<HashMap<String, String>>,
         tensors: Vec<(String, TensorInfo)>,
     ) -> Result<Self, SafeTensorError> {
-        let mut index_map = HashMap::new();
+        let mut index_map = HashMap::with_capacity(tensors.len());
 
         let tensors: Vec<_> = tensors
             .into_iter()
@@ -938,7 +938,7 @@ mod tests {
             .sum::<usize>()
             * dtype.size(); // 4
         let all_data = vec![0; n];
-        let mut metadata: HashMap<String, TensorView> = HashMap::new();
+        let mut metadata = HashMap::with_capacity(tensors_desc.len());
         let mut offset = 0;
         for (name, shape) in tensors_desc {
             let n: usize = shape.iter().product();