前言

我在刚接触的pytorch的时候，只有一台破笔记本，学到CNN的时候，需要用显卡训练模型，那时的我，兜比脸干净，此生头一次感觉到贫穷限制了我对知识的追求。

再回首恍然如梦，尝试垃圾卡一样可以训模型，我命由我不由天。

我的思路是这样：

1. 首先我们需要一个跨平台，支持多种显卡驱动的，统一的api框架。

2. 然后还需要一个能够在这个api之上，训练任意模型的媒介

3. 最后把我们的pytorch模型放在这个媒介上训练

1. wgpu

wgpu框架，是一个跨平台、安全、纯 Rust 的图形 API。它可以运行在 Vulkan、Metal、D3D12 和 OpenGL ，以及 wasm 上的 WebGL2 和 WebGPU 之上。

它的API基于WebGPU 标准实现。

看一下它支持的平台：

API	Windows	Linux/Android	macOS/iOS	Web (wasm)
Vulkan	✅	✅	🌋
Metal			✅
DX12	✅
OpenGL	🆗 (GL 3.3+)	🆗 (GL ES 3.0+)	📐	🆗 (WebGL2)
WebGPU				✅

✅ = 支持
🆗 = 低级支持
📐 = 需要ANGLE翻译层（仅限 GL ES 3.0）
🌋 = 需要MoltenVK翻译层
🛠️ = 不受支持，但欢迎贡献

可以看到他的跨平台，各种驱动的支持力度，都非常强，并且作为核心被应用在Firefox和Deno中，是我们的不二首选。

2. burn

burn是一个使用 Rust 构建的全新综合动态深度学习框架，
以极高的灵活性、计算效率和可移植性为主要目标。

作为中间的媒介，burn有两个特点是我选择的理由

• 对各种设备的兼容，这是burn最独特的一点
• 对于移植性做的特别好，是作为媒介最基本的特性

当然，作为一个新兴起的框架，burn还有很多待改进的地方，比如对新算法的支持，对灵活数据结构的支持，不过就工程角度上讲，burn是满足大部分需要的。

代码实战

我们这里用手写数字识别的例子，演示如何用我们笔记本的显卡训练。

模型代码，我们这里就用torch官方的mnist的代码

看一下我电脑的显卡： Intel UHD Graphics 630 1536 MB，标准的垃圾卡

1. pytorch模型与导出

我这里简单贴一下模型：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 3)
        self.conv2 = nn.Conv2d(8, 16, 3)
        self.conv3 = nn.Conv2d(16, 24, 3)
        self.norm1 = nn.BatchNorm2d(24)
        self.dropout1 = nn.Dropout(0.3)
        self.fc1 = nn.Linear(24 * 22 * 22, 32)
        self.fc2 = nn.Linear(32, 10)
        self.norm2 = nn.BatchNorm1d(10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = self.norm1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout1(x)
        x = self.fc2(x)
        x = self.norm2(x)
        output = F.log_softmax(x, dim=1)
        return output

1.1 state_dict方式导出

第一种方式：我们可以通过state_dict导出模型参数，如下代码，我们能导出mnist.pt的文件。

if args.save_model:
    torch.save(model.state_dict(), "mnist.pt")

然后在burn中，声明同样的模型：

#[derive(Module, Debug)]
pub struct Model<B: Backend> {
    conv1: Conv2d<B>,
    conv2: Conv2d<B>,
    conv3: Conv2d<B>,
    norm1: BatchNorm<B, 2>,
    fc1: Linear<B>,
    fc2: Linear<B>,
    norm2: BatchNorm<B, 0>,
    phantom: core::marker::PhantomData<B>,
}
...
    pub fn forward(&self, input1: Tensor<B, 4>) -> Tensor<B, 2> {
        let conv1_out1 = self.conv1.forward(input1);
        let relu1_out1 = relu(conv1_out1);
        let conv2_out1 = self.conv2.forward(relu1_out1);
        let relu2_out1 = relu(conv2_out1);
        let conv3_out1 = self.conv3.forward(relu2_out1);
        let relu3_out1 = relu(conv3_out1);
        let norm1_out1 = self.norm1.forward(relu3_out1);
        let flatten1_out1 = norm1_out1.flatten(1, 3);
        let fc1_out1 = self.fc1.forward(flatten1_out1);
        let relu4_out1 = relu(fc1_out1);
        let fc2_out1 = self.fc2.forward(relu4_out1);
        let norm2_out1 = self.norm2.forward(fc2_out1);
        log_softmax(norm2_out1, 1)
    }

最后将mnist加载到我们的模型中来。

    // 加载mnist.pt文件
    let record = NamedMpkFileRecorder::<FullPrecisionSettings>::default()
        .load(Path::new(OUT_DIR).into(), &device)
        .expect("Failed to decode state");

    // 创建模型，并把文件中的参数加载到模型里
    let model: Model<Backend> = Model::init(&device).load_record(record);

这种方式，简直是脱裤子放p，xxxx

1.2 onnx 开放神经网络交换

上面的方式能够加载torch的模型，但是，简直是一坨s，我们肯定是不会用的。

思考：我们的目标是星辰大海，怎么能局限在torch上那，既然要做开发能力，那我们索性就做到底，直接对onnx模型进行训练

ONNX（Open Neural Network Exchange）：是一套表示深度神经网络模型的开放格式，规范了 AI 模型交换标准，使 AI 模型可以在不同框架和环境下交互使用。

我们将pytorch中的模型，用如下代码，导出onnx格式。最终得到一个mnist.onnx文件。

    if args.export_onnx:
        dummy_input = torch.randn(1, 1, 28, 28, device=device)
        torch.onnx.export(model, dummy_input, "mnist.onnx",
                          verbose=True, opset_version=16)

根据这个文件，我们可以直接生成我们的模型

在rust项目中的build.rs中，构建mnist.onnx成为一个mnist.rs，构建包使用burn-import

ModelGen::new()
    .input("./model/mnist.onnx")
    .out_dir("./model/")
    .run_from_script();

注意： 在某些情况下，构建地址会有问题，比如workspace模式下，并不是构建在当前目录，然后你可能无法找到mnist.rs的地址。

解决方法： 把构建日志打印出来，命令cargo build -vv，如下图，可以看到文件的输出路径。

看一下自动构建的模型代码：

#[derive(Module, Debug)]
pub struct Model<B: Backend> {
    conv2d1: Conv2d<B>,
    conv2d2: Conv2d<B>,
    conv2d3: Conv2d<B>,
    batchnormalization1: BatchNorm<B, 2>,
    linear1: Linear<B>,
    linear2: Linear<B>,
    batchnormalization2: BatchNorm<B, 0>,
    phantom: core::marker::PhantomData<B>,
}
...

#[allow(clippy::let_and_return, clippy::approx_constant)]
pub fn forward(&self, input1: Tensor<B, 4>) -> Tensor<B, 2> {
    let conv2d1_out1 = self.conv2d1.forward(input1);
    let relu1_out1 = burn::tensor::activation::relu(conv2d1_out1);
    let conv2d2_out1 = self.conv2d2.forward(relu1_out1);
    let relu2_out1 = burn::tensor::activation::relu(conv2d2_out1);
    let conv2d3_out1 = self.conv2d3.forward(relu2_out1);
    let relu3_out1 = burn::tensor::activation::relu(conv2d3_out1);
    let batchnormalization1_out1 = self.batchnormalization1.forward(relu3_out1);
    let flatten1_out1 = batchnormalization1_out1.flatten(1, 3);
    let linear1_out1 = self.linear1.forward(flatten1_out1);
    let relu4_out1 = burn::tensor::activation::relu(linear1_out1);
    let linear2_out1 = self.linear2.forward(relu4_out1);
    let batchnormalization2_out1 = self.batchnormalization2.forward(linear2_out1);
    let logsoftmax1_out1 = burn::tensor::activation::log_softmax(batchnormalization2_out1, 1);
    logsoftmax1_out1
}

当然，我们也可以直接构建到我们的二进制程序里，

ModelGen::new()
    .input("./model/mnist.onnx")
    .out_dir("./model/")
    .record_type(RecordType::Bincode) //类型为bin
    .embed_states(true)
    .run_from_script();

当这种情况下，我们无法用传统的mod xxx的方式引用模型，需要用如下代码，引入构建后的mnist.rs。

pub mod mnist {
    include!(concat!(env!("OUT_DIR"), "./model/mnist.rs"));
}

2. 验证构建的模型

我们从pytorch导出的onnx模型，除了模型本身，还是具有参数的，实际在torch的例子中，导出的是一个训练好的卷积网络。

我们在自动构建之后，除了mnist.rs，还有mnist.bin和mnist.mpk,里面存放了模型参数等信息。

也就是说，我们可以直接将这个模型加载进来直接用。

加载模型：

let model: Model<Backend> = Model::default();

如上创建模型，会默认加载参数，如下代码，从./model/mnist目录中加载bin和mpk

impl<B: Backend> Default for Model<B> {
    fn default() -> Self {
        Self :: from_file ("./model/mnist" , & Default :: default ())
    }
}

我们构建一个验证过程：指定一个mnist测试集中的一个图片，让这个模型来识别数字是多少？

代码传送门

用如下命令运行程序

cargo run -- 12

效果截图，可以看到能够精准预测。

3. 准备训练无参模型

上面验证了有参模型，通过构建是能够直接用的，但是，我们的目标是训练一个没调整过的模型。

3.1 实现数据加载

和torch的套路一样，burn同样需要构建dataset和dataloader，代码如下，我这里直接用了burn的MNISTDataset的数据集。主要功能就是数据加载 和 预处理。

#[derive(Clone)]
pub struct ClassificationBatcher<B: Backend> {
    device: B::Device,
}

#[derive(Clone, Debug)]
pub struct ClassificationBatch<B: Backend> {
    pub images: Tensor<B, 4>,
    pub targets: Tensor<B, 1, Int>,
}

impl<B: Backend> ClassificationBatcher<B> {
    pub fn new(device: B::Device) -> Self {
        Self {
            device,
        }
    }
}

impl<B: Backend> Batcher<MNISTItem, ClassificationBatch<B>> for ClassificationBatcher<B> {
    fn batch(&self, items: Vec<MNISTItem>) -> ClassificationBatch<B> {
        let targets = items
            .iter()
            .map(|item| {
                Tensor::<B, 1, Int>::from_data(Data::from([(item.label as i64).elem()]), &self.device)
            })
            .collect();

        let images = items
            .into_iter()
            .map(|item| {
                let image_data = item.image.iter().copied().flatten().collect::<Vec<f32>>();
                let mut input: Tensor<B, 3> =
                    Tensor::from_floats(image_data.as_slice(), &self.device).reshape([1, 28, 28]);
                // Normalize the input
                input = ((input / 255) - 0.1307) / 0.3081;
                input
            })
            .collect();

        let images = Tensor::stack(images, 0);
        let targets = Tensor::cat(targets, 0);

        ClassificationBatch { images, targets }
    }
}

3.2 迭代和损失策略

• 我们这里使用经典的CrossEntropyLoss损失函数
• 优化器使用AdaGrad
• TrainStep是模型用于训练的trait，ValidStep是模型用于验证的trait，我们要为自动生成的Model，实现这两个trait，才能开始训练

impl<B: Backend> Model<B> {
    pub fn forward_classification(
        &self,
        images: Tensor<B, 4>,
        targets: Tensor<B, 1, Int>,
    ) -> ClassificationOutput<B> {
        let output = self.forward(images);
        let loss = CrossEntropyLossConfig::new()
            .init(&output.device())
            .forward(output.clone(), targets.clone());

        ClassificationOutput::new(loss, output, targets)
    }
}

impl<B: AutodiffBackend> TrainStep<ClassificationBatch<B>, ClassificationOutput<B>> for Model<B> {
    fn step(&self, batch: ClassificationBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
        let item = self.forward_classification(batch.images, batch.targets);

        TrainOutput::new(self, item.loss.backward(), item)
    }
}

impl<B: Backend> ValidStep<ClassificationBatch<B>, ClassificationOutput<B>> for Model<B> {
    fn step(&self, batch: ClassificationBatch<B>) -> ClassificationOutput<B> {
        self.forward_classification(batch.images, batch.targets)
    }
}

3.3 训练过程，和添加监控

• 下面是我们的整个训练过程，套路和torch一模一样
• 我们将训练用的配置，和最终的模型，保存到文件中
• 加入burn自带的监控，方便我们观察准确率和损失的变化。
• 我们这里用new方法加载默认参数的模型，也就是没训练过的模型，而不是default

pub fn train<B: AutodiffBackend>(config: TrainingConfig, device: B::Device) {
    std::fs::create_dir_all(ARTIFACT_DIR).ok();
    config
        .save(format!("{ARTIFACT_DIR}/config.json"))
        .expect("Config should be saved successfully");

    B::seed(config.seed);

    let dataset_train = MNISTDataset::train();
    let dataset_test = MNISTDataset::train();

    // Dataloaders
    let batcher_train = ClassificationBatcher::<B>::new(device.clone());
    let batcher_valid = ClassificationBatcher::<B::InnerBackend>::new(device.clone());

    let dataloader_train = DataLoaderBuilder::new(batcher_train)
        .batch_size(config.batch_size)
        .shuffle(config.seed)
        .num_workers(config.num_workers)
        .build(dataset_train);

    let dataloader_test = DataLoaderBuilder::new(batcher_valid)
        .batch_size(config.batch_size)
        .num_workers(config.num_workers)
        .build(dataset_test);

    // Learner config
    let learner = LearnerBuilder::new(ARTIFACT_DIR)
        .metric_train_numeric(AccuracyMetric::new())
        .metric_valid_numeric(AccuracyMetric::new())
        .metric_train_numeric(LossMetric::new())
        .metric_valid_numeric(LossMetric::new())
        .with_file_checkpointer(CompactRecorder::new())
        .devices(vec![device.clone()])
        .num_epochs(config.num_epochs)
        .build(
            Model::new(&device.clone()),
            config.optimizer.init(),
            config.learning_rate,
        );

    // Training
    let now = Instant::now();
    let model_trained = learner.fit(dataloader_train, dataloader_test);
    let elapsed = now.elapsed().as_secs();
    println!("Training completed in {}m{}s", (elapsed / 60), elapsed % 60);

    model_trained
        .save_file(format!("{ARTIFACT_DIR}/model"), &CompactRecorder::new())
        .expect("Trained model should be saved successfully");
}

4. 开始训练

正式训练之前，需要配置一下训练相关的参数，将这个cfg传入train即可。

let mut cfg = TrainingConfig::new(AdaGradConfig::new());
cfg.num_workers = 4;
cfg.num_epochs = 8;
cfg.batch_size = 1000;
cfg.learning_rate = 1.0;

4.1 CPU

我们先用cpu尝试一下：

train::<Autodiff<NdArray>>(cfg, NdArrayDevice::Cpu);

run起程序，

风扇开始呼呼转动，窗口直接黑屏卡死，cpu监控拉满：

4.2 GPU

接下来用我们的笔记本自带的垃圾卡跑一下（一般的a卡和i卡都可以跑）

train::<Autodiff<Wgpu>>(cfg, WgpuDevice::default());

看一下监控的变化：

• 预计用时37分分钟
• 准确率增加明显

再看一下损失变化情况，降低的也很好

看一下显卡的使用情况，a卡已经被利用起来了

最后看一下cpu的负载，有增加，但不多

4.3 libtorch

假如你很富有，有n卡，或者m1的mac，那么你同样可以用这种方式训练。

比如用m1芯片：

train::<Autodiff<LibTorch>>(cfg,LibTorchDevice::Mps);

或者CUDA

train::<Autodiff<LibTorch>>(cfg,LibTorchDevice::Cuda(0));

我这里就不测试了

尾语

burn的能力不仅限于上面的种种，它允许自定义设备的接入。但是目前还处在较为初期的阶段，像我们pytorch例子中的nll_loss和Adadelta都是不支持的。

前几年挖矿，和最近的ai的尽头是电力和算力，都在大力拉升显卡的价格，国内显卡和好的芯片价格高得离谱，还要处处被卡脖子，哎~

最后祝愿每个点赞收藏的帅哥美女都用上好卡。