企业官网建设流程全解析

本篇指南围绕 Koharu 开源仓库（ML 驱动的漫画翻译器，核心用 Rust 编写）中负责深度学习计算的 koharu-torch crate 展开，介绍它如何以 tch-rs 风格封装 PyTorch 的 C++ API（LibTorch），并通过动态加载 shim 的方式在运行时接入原生库。读完本文，你将掌握：Tensor 基础运算、用nn::VarStore与优化器完成梯度下降训练、编写简单神经网络（MNIST 分类）、加载预训练视觉模型，以及通过 SafeTensors 格式在 Python 与 Rust 之间迁移模型权重，并了解 Koharu 运行时是如何为这些能力提供 LibTorch 原生依赖的。

项目定位：面向 Rust 的 LibTorch 绑定

koharu-torch是经典tch-rs的衍生实现。它的设计目标是提供围绕 PyTorch C++ API（即 libtorch）的薄封装，尽量贴近原始 C++ API 的风格，以便在此之上继续开发更符合 Rust 习惯的高层绑定。从 crates/koharu-torch/Cargo.toml 可以看到，该 crate 依赖底层 FFI 绑定koharu-torch-sys，并引入了safetensors（SafeTensors 读写）、half（半精度浮点）、image（图像处理）、ndarray、rand、zip等支撑库。

在入口文件 crates/koharu-torch/src/lib.rs 中，顶层 API 一览无余：

• 张量体系：Tensor、Shape、IndexOp、NewAxis、Reduction，以及autocast、no_grad、no_grad_guard、with_grad等自动求导开关；
• 设备与数据类型：Device（CPU / CUDA / MPS / Vulkan）、Cuda、Kind；
• 神经网络高层库：nn模块；
• 预训练视觉模型与数据集：vision模块；
• TorchScript 支持：jit模块中的CModule、IValue、TrainableCModule；
• 线程与随机种子控制：set_num_threads、get_num_threads、manual_seed等。

整个 crate 采用#[macro_use] extern crate lazy_static配合底层 FFI 调用，所有异常通过统一错误类型TchError收敛（见 crates/koharu-torch/src/lib.rs）。

环境准备：LibTorch v2.13.0 的动态加载机制

README 的 Getting Started 一节明确了本项目与上游 tch-rs 的最大区别：本 fork 动态加载针对 LibTorch v2.13.0 构建的koharu-torch库。Koharu 的运行时负责安装配套的原生归档，并在张量 API 被使用之前激活 LibTorch 及其捆绑的 shim。

这一机制在源码中有完整印证。运行时侧的包描述见 crates/koharu-runtime/src/runtime/packages/torch.rs：

• 版本常量RELEASE = "v2.13.0.7"，与 README 所述 LibTorch v2.13.0 对应；
• Torch枚举划分三种后端：Cpu（macOS 上序列化为metal）、Cuda、Rocm（HIP），覆盖 CPU、NVIDIA CUDA、AMD ROCm 三大计算后端；
• 每个后端按平台声明了一串必须存在的动态库清单，例如 Linux CPU 需要libgomp.so.1, libc10.so, libtorch_global_deps.so, libtorch_cpu.so, libtorch.so, libkoharu-torch.so，Windows 对应libiomp5md.dll, c10.dll, torch_global_deps.dll, torch_cpu.dll, torch.dll, koharu-torch.dll，macOS 则是libtorch.dylib, libtorch_global_deps.dylib, libtorch_cpu.dylib, libc10.dylib, libkoharu-torch.dylib；
• install()会检查本地归档是否已完整（complete方法逐文件校验），否则从发布通道下载与目标三元组匹配的.tar.gz归档，并只抽取动态库文件（**/*.dll、**/*.dylib、**/*.so、**/*.so.*）；
• 硬件发现逻辑discover()会优先选择 Metal 支持（macOS）或按 Linux 架构/GPU 能力决定 CPU 还是 CUDA/ROCm 包。

在 FFI 层，crates/koharu-torch-sys/src/lib.rs 用OnceLock缓存函数指针，通过__koharu_bindgen_load在首次调用时按符号名动态解析（例如at_autocast_is_enabled），函数签名则来自构建期生成的torch_api.rs。构建脚本 crates/koharu-torch-sys/build.rs 借助koharu-bindgen从 crates/koharu-torch-sys/libtch/torch_api.h 与torch_api_generated.h生成绑定：torch::Tensor *、torch::Scalar *、torch::optim::Optimizer *、torch::jit::script::Module *、torch::jit::IValue *分别映射为tensor、scalar、optimizer、torch_module、ivalue不透明指针，所有调用包裹在PROTECT宏中捕获 C++ 异常并写入线程本地错误串。这也是 README 中"动态加载 + 运行时安装原生归档"的底层实现。

Windows 专项注意事项

README 中特别提醒：依据 PyTorch 官方文档，Windows 上 Debug 与 Release 构建不兼容 ABI，如果使用了与编译产物不匹配的 libtorch 版本，可能导致段错误。

同时建议使用MSVC Rust 工具链（例如通过 rustup 安装stable-x86_64-pc-windows-msvc），而非基于 MinGW 的工具链，因为 PyTorch 与 MinGW 存在兼容性问题。在 crates/koharu-runtime/src/runtime/packages/torch.rs 中也能看到，Windows 只支持x86_64-pc-windows-msvc目标，与 README 的建议一致。

基础张量操作

Tensor 类型是对 PyTorch Tensor 的封装。README 给出的最小示例：

use tch::Tensor; fn main() { let t = Tensor::from_slice(&[3, 1, 4, 1, 5]); let t = t * 2; t.print(); }

在 koharu-torch 中同样可用（crate 名替换为koharu_torch）。从源码看，Tensor 持有底层 C++ 指针*mut C_tensor（见 crates/koharu-torch/src/wrappers/tensor.rs），并通过unsafe_torch!宏调用at_new_tensor、at_shallow_clone、at_shape、at_stride、at_copy_data等 C API（声明见 crates/koharu-torch-sys/libtch/torch_api.h）。常用方法包括：

• dim()/size()/size1()/size2()：维度与形状查询，底层对应at_dim、at_shape；
• from_slice、from_data_size：从 Rust 切片/原始数据构造张量，对应at_tensor_of_data/at_tensor_of_blob；
• view/f_view：重塑张量，Shapetrait 已为[i64; N]、元组、&[i64]、i64等类型实现（见 crates/koharu-torch/src/tensor/mod.rs），所以可以直接写t.view(&[2, 5])；
• pow_tensor_scalar、exp、sum等运算方法由自动生成的tensor_generated.rs提供。

张量元素类型由Kind枚举管理（见 crates/koharu-torch/src/wrappers/kind.rs），覆盖Uint8、Int8、Int16、Int、Int64、Half、Float、Double、Bool、BFloat16、QInt8、Float8e5m2等类型，其枚举值与 PyTorch 的c10/core/ScalarType.h保持一致；elt_size_in_bytes()可查询每个元素的字节数。FLOAT_CPU、INT64_CPU、FLOAT_CUDA等预定义常量（如(Kind::Float, Device::Cpu)）可作便捷组合。

用梯度下降训练模型

PyTorch 为大多数张量运算提供自动微分，常用于梯度下降训练。训练在nn::VarStore创建的变量上进行，优化器则通过nn::Sgd构建。README 的完整示例：

use tch::nn::{Module, OptimizerConfig}; use tch::{kind, nn, Device, Tensor}; fn my_module(p: nn::Path, dim: i64) -> impl nn::Module { let x1 = p.zeros("x1", &[dim]); let x2 = p.zeros("x2", &[dim]); nn::func(move |xs| xs * &x1 + xs.exp() * &x2) } fn gradient_descent() { let vs = nn::VarStore::new(Device::Cpu); let my_module = my_module(vs.root(), 7); let mut opt = nn::Sgd::default().build(&vs, 1e-2).unwrap(); for _idx in 1..50 { // Dummy mini-batches made of zeros. let xs = Tensor::zeros(&[7], kind::FLOAT_CPU); let ys = Tensor::zeros(&[7], kind::FLOAT_CPU); let loss = (my_module.forward(&xs) - ys).pow_tensor_scalar(2).sum(kind::Kind::Float); opt.backward_step(&loss); } }

这里的my_module使用两个初始值为 0 的变量x1、x2，前向传播计算xs * x1 + exp(xs) * x2。每个训练迭代包含三步：对 mini-batch 执行前向传播 → 计算模型输出与 ground truth 的均方误差作为 loss → 执行优化步（计算梯度并更新VarStore中的变量）。

源码侧的关键实现：

• crates/koharu-torch/src/nn/var_store.rs：VarStore内部维护Variables（named_variables: HashMap<String, Tensor>与trainable_variables: Vec<Var>），用Arc<Mutex<...>>保护，并固定单一设备；Path负责变量命名，路径元素以.分隔，支持vs.root()、vs / "layer1"这类层级路径；
• crates/koharu-torch/src/nn/optimizer.rs：OptimizerConfig::build把VarStore中所有可训练变量按分组加入 C++ 优化器（opt.add_parameters），返回Optimizer；随后backward_step(&loss)完成反向传播与参数更新；
• 除Sgd（默认momentum=0, dampening=0, wd=0, nesterov=false）外，还提供Adam（默认beta1=0.9, beta2=0.999, eps=1e-8）、AdamW、RmsProp，均有default()与便捷构造函数sgd(...)、adam(...)。

编写一个简单的神经网络

nnAPI 可用于搭建网络结构。README 给出了带一个隐藏层、在 MNIST 上用 Adam 优化器训练的模型：

use anyhow::Result; use tch::{nn, nn::Module, nn::OptimizerConfig, Device}; const IMAGE_DIM: i64 = 784; const HIDDEN_NODES: i64 = 128; const LABELS: i64 = 10; fn net(vs: &nn::Path) -> impl Module { nn::seq() .add(nn::linear( vs / "layer1", IMAGE_DIM, HIDDEN_NODES, Default::default(), )) .add_fn(|xs| xs.relu()) .add(nn::linear(vs, HIDDEN_NODES, LABELS, Default::default())) } pub fn run() -> Result<()> { let m = tch::vision::mnist::load_dir("data")?; let vs = nn::VarStore::new(Device::Cpu); let net = net(&vs.root()); let mut opt = nn::Adam::default().build(&vs, 1e-3)?; for epoch in 1..200 { let loss = net .forward(&m.train_images) .cross_entropy_for_logits(&m.train_labels); opt.backward_step(&loss); let test_accuracy = net .forward(&m.test_images) .accuracy_for_logits(&m.test_labels); println!( "epoch: {:4} train loss: {:8.5} test acc: {:5.2}%", epoch, f64::from(&loss), 100. * f64::from(&test_accuracy), ); } Ok(()) }

该示例用到三块能力：

1. 序列式网络：nn::seq().add(...).add_fn(|xs| xs.relu())，在 crates/koharu-torch/src/nn/sequential.rs 中实现；
2. 数据集：tch::vision::mnist::load_dir("data")加载 MNIST 训练/测试图像与标签。vision模块还内置 CIFAR-10 等数据集以及 AlexNet、ResNet、DenseNet、VGG、SqueezeNet、MobileNet、EfficientNet、ConvMixer、DINOv2 等模型定义（见 crates/koharu-torch/src/vision/mod.rs）；
3. 训练指标：cross_entropy_for_logits计算交叉熵损失，accuracy_for_logits计算分类准确率，f64::from(&loss)把张量转为标量。nn::Module/ModuleTtrait 在 crates/koharu-torch/src/nn/module.rs 中定义，CModule、TrainableCModule也都实现了这两个 trait，因此 TorchScript 模块可以直接参与前向与训练流程。

使用预训练视觉模型

README 的pretrained-models示例演示了如何在图片上使用预训练模型（权重从 PyTorch 实现提取，如 resnet18/resnet34），运行命令为：

cargo run --example pretrained-models -- resnet18.ot tiger.jpg

输出该图片在 ImageNet 上的前 5 个类别及概率。核心代码逻辑：

// First the image is loaded and resized to 224x224. let image = imagenet::load_image_and_resize(image_file)?; // A variable store is created to hold the model parameters. let vs = tch::nn::VarStore::new(tch::Device::Cpu); // Then the model is built on this variable store, and the weights are loaded. let resnet18 = tch::vision::resnet::resnet18(vs.root(), imagenet::CLASS_COUNT); vs.load(weight_file)?; // Apply the forward pass of the model to get the logits and convert them // to probabilities via a softmax. let output = resnet18 .forward_t(&image.unsqueeze(0), /*train=*/ false) .softmax(-1); // Finally print the top 5 categories and their associated probabilities. for (probability, class) in imagenet::top(&output, 5).iter() { println!("{:50} {:5.2}%", class, 100.0 * probability) }

流程是：加载并缩放图片至 224×224 → 创建VarStore并在其上构建resnet18(vs.root(), imagenet::CLASS_COUNT)→vs.load(weight_file)装载权重 → 以forward_t(..., false)推理（train 标志为 false）得到 logits →softmax(-1)转为概率 →imagenet::top取 top-5。相关的resnet模型定义、imagenet辅助函数都在vision模块中（crates/koharu-torch/src/vision/resnet.rs、crates/koharu-torch/src/vision/imagenet.rs），图像加载与缩放在 crates/koharu-torch/src/wrappers/image.rs 中通过at_load_image、at_resize_image完成。

用 SafeTensors 导入/导出 PyTorch 权重

SafeTensors是 HuggingFace 提出的张量存储格式：不依赖 Python 的pickle模块，张量不与保存时的类及目录结构绑定，同时支持零拷贝读取——读取文件所需内存不超过文件本身大小。

安装safetensors

在 Python 侧安装：

pip install safetensors

在 PyTorch 中导出权重

import torchvision from safetensors import torch as stt model = torchvision.models.resnet18(pretrained=True) stt.save_file(model.state_dict(), 'resnet18.safetensors')

注意：导出文件名必须以.safetensors后缀结尾，tch 才能正确解码。

在tch中导入权重

use anyhow::Result; use tch::{ Device, Kind, nn::VarStore, vision::{ imagenet, resnet::resnet18, } }; fn main() -> Result<()> { // Create the model and load the pre-trained weights let mut vs = VarStore::new(Device::cuda_if_available()); let model = resnet18(&vs.root(), 1000); vs.load("resnet18.safetensors")?; // Load the image file and resize it to the usual imagenet dimension of 224x224. let image = imagenet::load_image_and_resize224("dog.jpg")? .to_device(vs.device()); // Apply the forward pass of the model to get the logits let output = image .unsqueeze(0) .apply_t(&model, false) .softmax(-1, Kind::Float); // Print the top 5 categories for this image. for (probability, class) in imagenet::top(&output, 5).iter() { println!("{:50} {:5.2}%", class, 100.0 * probability) } Ok(()) }

注意这里VarStore::new(Device::cuda_if_available())会优先选择 CUDA 设备；Device::cuda_if_available()的实现见 crates/koharu-torch/src/wrappers/device.rs，通过Cuda::is_available()判断是否回退到 CPU。Kind::Float显式指定 softmax 输出类型。

源码级的 SafeTensors 支持

koharu-torch 对 SafeTensors 的支持远不止vs.load这一条路径。crates/koharu-torch/src/tensor/safetensors.rs 实现了完整的读写能力：

• Tensor::read_safetensors(path)：读取文件，返回Vec<(String, Tensor)>命名张量列表；
• Tensor::write_safetensors(&[(name, tensor)], path)：把命名张量序列化写入文件（serialize_to_file）；
• VarStore::read_safetensors(path)：按VarStore中已有变量名逐项匹配文件中的张量并复制，缺失的张量会报错；
• VarStore::fill_safetensors(path)：宽松版本，只填充名字匹配的变量；
• Kind与 SafeTensorsDtype之间实现了双向TryFrom转换（F32↔Float、F16↔Half、BF16↔BFloat16、I64↔Int64等，见 crates/koharu-torch/src/tensor/safetensors.rs）；
• 写出前会校验：稀疏张量与非连续张量不允许保存（is_sparse/is_contiguous检查）；
• 模块内自带parse单元测试（crates/koharu-torch/src/tensor/safetensors.rs），验证各类Kind/Dtype相互映射的正确性。

在 Koharu 中落地：koharu-ml 的实际用法

koharu-torch 并非孤立存在——它是整个 Koharu ML 推理栈的计算底座。从源码搜索可以看到，多个模型模块直接依赖koharu_torch：

• crates/koharu-ml/src/aot_inpainting/model.rs 的Model::new先创建nn::VarStore::new(device)，用vs.root()下的层级路径（如root / "head" / 0）搭建 GatedWSConv 卷积网络，随后vs.freeze()冻结参数并调用self.vs.load(path)?从model.safetensors加载权重；
• crates/koharu-ml/src/backend.rs 的set_precision会根据硬件能力自动切换精度：CUDA 设备计算能力 ≥ 80（或 ROCm 的 gfx908/gfx90a/gfx11/gfx12 系列）时使用BFloat16，否则回退Float——这正是Kind::BFloat16在实际推理中的典型应用；
• 推理时普遍使用koharu_torch::no_grad(...)包裹前向调用（如 crates/koharu-ml/src/aot_inpainting/mod.rs），避免梯度计算开销。

这说明 koharu-torch 的nn、VarStore、Kind、Device、no_grad等 API 已经构成了生产级模型加载与推理的基础设施。

FAQ：动态库加载报错排查

README 给出了最常见的运行期问题：生成二进制后找不到共享库，例如：

error while loading shared libraries: libtorch_cpu.so: cannot open shared object file: No such file or directory

解决方案是在 shell 配置中把 libtorch 安装目录加入动态库搜索路径（/path/to/libtorch替换为实际路径）：

# For Linux export LD_LIBRARY_PATH=/path/to/libtorch/lib:$LD_LIBRARY_PATH # For macOS export DYLD_LIBRARY_PATH=/path/to/libtorch/lib:$DYLD_LIBRARY_PATH

在 Koharu 的实际架构中，这一职责已由运行时接管：koharu-runtime会先安装并校验与平台匹配的原生归档（见上文"动态加载机制"一节），因此在应用内部通常无需手动配置环境变量。

更多示例方向

README 列举了丰富的进阶用例方向（均对应上游 tch-rs 生态的经典示例）：字符级语言建模（char-RNN）、基于预训练 VGG-16 的神经风格迁移、CIFAR-10 上的 ResNet 训练、用 TorchScript JIT 部署 Python 训练模型、强化学习（OpenAI Gym 上的策略梯度与 A2C）、小数据集上的迁移学习、简化版 GPT，以及参照 HuggingFace diffusers 路线实现的 Stable Diffusion。在 koharu-torch 中，jit模块提供了CModule/TrainableCModule与IValue支持（crates/koharu-torch/src/wrappers/jit.rs），可用于加载.pt/TorchScript 模块并与 Tensor 互操作。

许可证

koharu-torch与上游 tch-rs 一致，采用 MIT 与 Apache 2.0 双许可证，任选其一。仓库根目录与 crate 目录下分别提供了 LICENSE-APACHE 与 LICENSE-MIT 全文。

小结

koharu-torch 为 Rust 生态提供了一套贴近 PyTorch C++ API 的完整深度学习工具链：从Tensor/Kind/Device基础张量体系，到nn模块的层、序列、VarStore变量存储与多种优化器，再到vision模块的预训练模型与数据集、jit模块的 TorchScript 加载，以及 SafeTensors 格式的权重迁移。结合 crates/koharu-runtime/src/runtime/packages/torch.rs 的运行时原生依赖管理与 crates/koharu-ml 中各模型模块的实际调用，你可以直接参考本指南中的示例，在 Rust 项目中完成从张量运算、模型训练、预训练权重加载到生产推理的完整闭环。