tch-serde: 使用 serde 序列化和反序列化 tch-rs 类型

这个 crate 为 tch-rs 常见类型提供了 {ser,de} ialization 方法。

docs.rs | crates.io

使用方法

例如，为 tensor 字段启用序列化，添加以下注解：#[serde(with = "tch_serde::serde_tensor")]

use tch::{Device, Kind, Tensor};

#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct Example {
    #[serde(with = "tch_serde::serde_tensor")]
    tensor: Tensor,
    #[serde(with = "tch_serde::serde_kind")]
    kind: Kind,
    #[serde(with = "tch_serde::serde_device")]
    device: Device,
    #[serde(with = "tch_serde::serde_reduction")]
    reduction: Reduction,
}

fn main() {
    let example = Example {
        tensor: Tensor::randn(&[2, 3], (Kind::Float, Device::Cuda(0))),
        kind: Kind::Float,
        device: Device::Cpu,
        reduction: Reduction::Mean,
    };
    let text = serde_json::to_string_pretty(&example).unwrap();
    println!("{}", text);
}

许可证

MIT 许可证。请参阅 LICENSE 文件。