7 unstable releases (3 breaking)

0.4.1 Mar 26, 2024
0.4.0 Aug 21, 2023
0.3.0 Jul 18, 2023
0.2.2 Jun 11, 2023
0.1.0 May 9, 2023

#46 in Machine learning

Download history 22/week @ 2024-02-27 10/week @ 2024-03-05 7/week @ 2024-03-12 117/week @ 2024-03-26 62/week @ 2024-04-02 35/week @ 2024-04-09 37/week @ 2024-04-16 63/week @ 2024-04-23 159/week @ 2024-04-30 2107/week @ 2024-05-07 3701/week @ 2024-05-14 2359/week @ 2024-05-21 6029/week @ 2024-05-28 4568/week @ 2024-06-04 6071/week @ 2024-06-11

19,491 downloads per month

Apache-2.0

43KB
993 lines

dlpark

Github Actions Crates.io docs.rs

A pure Rust implementation of dmlc/dlpack.

Check example/with_pyo3 for usage.

This implementation focuses on transferring tensor from Rust to Python and vice versa.

It can also be used without pyo3 as a Rust library with default-features = false, check example/from_numpy.

Quick Start

We provide a simple example of how to transfer image::RgbImage to Python and torch.Tensor to Rust.

Full code is here.

Rust $\rightarrow$ Python

We have to implement some traits for a struct to be able to converted to PyObject

use std::ffi::c_void;
use dlpark::prelude::*;

struct PyRgbImage(image::RgbImage);

impl ToTensor for PyRgbImage {
    fn data_ptr(&self) -> *mut std::ffi::c_void {
        self.0.as_ptr() as *const c_void as *mut c_void
    }

    fn byte_offset(&self) -> u64 {
        0
    }


    fn device(&self) -> Device {
        Device::CPU
    }

    fn dtype(&self) -> DataType {
        DataType::U8
    }

    fn shape_and_strides(&self) -> ShapeAndStrides {
        ShapeAndStrides::new_contiguous_with_strides(
            &[self.0.height(), self.0.width(), 3].map(|x| x as i64),
        )
    }
}

Then we can return a ManagerCtx<PyRgbImage>

#[pyfunction]
fn read_image(filename: &str) -> ManagerCtx<PyRgbImage> {
    let img = image::open(filename).unwrap();
    let rgb_img = img.to_rgb8();
    ManagerCtx::new(PyRgbImage(rgb_img))
}

You can acess it in Python

import dlparkimg
from torch.utils.dlpack import to_dlpack, from_dlpack
import matplotlib.pyplot as plt

tensor = from_dlpack(dlparkimg.read_image("candy.jpg"))

print(tensor.shape)
plt.imshow(tensor.numpy())
plt.show()

If you want to convert it to numpy.ndarray , you can make a simple wrapper like this

import numpy as np
import dlparkimg

class FakeTensor:

    def __init__(self, x):
        self.x = x

    def __dlpack__(self):
        return self.x

arr = np.from_dlpack(FakeTensor(dlparkimg.read_image("candy.jpg")))

Python $\rightarrow$ Rust

ManagedTensor holds the memory of tensor and provide methods to access the tensor's attributes.

#[pyfunction]
fn write_image(filename: &str, tensor: ManagedTensor) {
    let buf = tensor.as_slice::<u8>();

    let rgb_img = image::ImageBuffer::<Rgb<u8>, _>::from_raw(
        tensor.shape()[1] as u32,
        tensor.shape()[0] as u32,
        buf,
    )
    .unwrap();

    rgb_img.save(filename).unwrap();
}

And you can call it in Python

import dlparkimg
from torch.utils.dlpack import to_dlpack, from_dlpack

bgr_img = tensor[..., [2, 1, 0]] # [H, W, C=3]
dlparkimg.write_image('bgr.jpg', to_dlpack(bgr_img))

Dependencies

~0–5.5MB
~15K SLoC