摆脱CUDA，用垃圾显卡训练pytorch搭建的神经网络 前言 我在刚接触的pytorch的时候，只有一台破笔记本，学到CNN的时候，需要用显卡训练模型，那时的我，兜比脸干净，此生头一次感觉到贫穷限制了我对知识的追求。 再回首恍然如梦，尝试垃圾卡一样可以训模型，我命由我不由天。 我的思路是这样： 首先我们需要一个跨平台，支持多种显卡驱动的，统一的api框架。 然后还需要一个能够在这个api之上，训练任意模型的媒介 最后把我们的pytorch模型放在这个媒介上训练 1. wgpu wgpu框架，是一个跨平台、安全、纯 Rust 的图形 API。它可以运行在 Vulkan、Metal、D3D12 和 OpenGL ，以及 wasm 上的 WebGL2 和 WebGPU 之上。 它的API基于WebGPU 标准实现。 看一下它支持的平台： APIWindowsLinux/AndroidmacOS/iOSWeb (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模型与导出我这里简单贴一下模型： classNet(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) defforward(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) returnoutput 1.1 state_dict方式导出第一种方式：我们可以通过state_dict导出模型参数，如下代码，我们能导出mnist.pt的文件。 ifargs.save_model: torch.save(model.state_dict(),"mnist.pt") 然后在burn中，声明同样的模型： #[derive(Module,Debug)] pubstructModelB:Backend{ conv1:Conv2d, conv2:Conv2d, conv3:Conv2d, norm1:BatchNormB,2, fc1:Linear, fc2:Linear, norm2:BatchNormB,0, phantom:core::marker::PhantomData, } ... pubfnforward(&self,input1:TensorB,4)-TensorB,2{ letconv1_out1=self.conv1.forward(input1); letrelu1_out1=relu(conv1_out1); letconv2_out1=self.conv2.forward(relu1_out1); letrelu2_out1=relu(conv2_out1); letconv3_out1=self.conv3.forward(relu2_out1); letrelu3_out1=relu(conv3_out1); letnorm1_out1=self.norm1.forward(relu3_out1); letflatten1_out1=norm1_out1.flatten(1,3); letfc1_out1=self.fc1.forward(flatten1_out1); letrelu4_out1=relu(fc1_out1); letfc2_out1=self.fc2.forward(relu4_out1); letnorm2_out1=self.norm2.forward(fc2_out1); log_softmax(norm2_out1,1) } 最后将mnist加载到我们的模型中来。 //加载mnist.pt文件 letrecord=NamedMpkFileRecorder::FullPrecisionSettings::default() .load(Path::new(OUT_DIR).into(),&device) .expect("Failedtodecodestate"); //创建模型，并把文件中的参数加载到模型里 letmodel:ModelBackend=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文件。 ifargs.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，如下图，可以看到文件的输出路径。 image.png看一下自动构建的模型代码： #[derive(Module,Debug)] pubstructModelB:Backend{ conv2d1:Conv2d, conv2d2:Conv2d, conv2d3:Conv2d, batchnormalization1:BatchNormB,2, linear1:Linear, linear2:Linear, batchnormalization2:BatchNormB,0, phantom:core::marker::PhantomData, } ... #[allow(clippy::let_and_return,clippy::approx_constant)] pubfnforward(&self,input1:TensorB,4)-TensorB,2{ letconv2d1_out1=self.conv2d1.forward(input1); letrelu1_out1=burn::tensor::activation::relu(conv2d1_out1); letconv2d2_out1=self.conv2d2.forward(relu1_out1); letrelu2_out1=burn::tensor::activation::relu(conv2d2_out1); letconv2d3_out1=self.conv2d3.forward(relu2_out1); letrelu3_out1=burn::tensor::activation::relu(conv2d3_out1); letbatchnormalization1_out1=self.batchnormalization1.forward(relu3_out1); letflatten1_out1=batchnormalization1_out1.flatten(1,3); letlinear1_out1=self.linear1.forward(flatten1_out1); letrelu4_out1=burn::tensor::activation::relu(linear1_out1); letlinear2_out1=self.linear2.forward(relu4_out1); letbatchnormalization2_out1=self.batchnormalization2.forward(linear2_out1); letlogsoftmax1_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。 pubmodmnist{ include!(concat!(env!("OUT_DIR"),"./model/mnist.rs")); } 2. 验证构建的模型我们从pytorch导出的onnx模型，除了模型本身，还是具有参数的，实际在torch的例子中，导出的是一个训练好的卷积网络。 我们在自动构建之后，除了mnist.rs，还有mnist.bin和mnist.mpk,里面存放了模型参数等信息。 也就是说，我们可以直接将这个模型加载进来直接用。 加载模型： letmodel:ModelBackend=Model::default(); 如上创建模型，会默认加载参数，如下代码，从./model/mnist目录中加载bin和mpk implB:BackendDefaultforModel{ fndefault()-Self{ Self::from_file("./model/mnist",Default::default()) } } 我们构建一个验证过程：指定一个mnist测试集中的一个图片，让这个模型来识别数字是多少？ 代码传送门 用如下命令运行程序 cargorun--12 效果截图，可以看到能够精准预测。 3. 准备训练无参模型上面验证了有参模型，通过构建是能够直接用的，但是，我们的目标是训练一个没调整过的模型。 3.1 实现数据加载和torch的套路一样，burn同样需要构建dataset和dataloader，代码如下，我这里直接用了burn的MNISTDataset的数据集。主要功能就是数据加载 和 预处理。 #[derive(Clone)] pubstructClassificationBatcherB:Backend{ device:B::Device, } #[derive(Clone,Debug)] pubstructClassificationBatchB:Backend{ pubimages:TensorB,4, pubtargets:TensorB,1,Int, } implB:BackendClassificationBatcher{ pubfnnew(device:B::Device)-Self{ Self{ device, } } } implB:BackendBatcherMNISTItem,ClassificationBatchforClassificationBatcher{ fnbatch(&self,items:VecMNISTItem)-ClassificationBatch{ lettargets=items .iter() .map(|item|{ Tensor::B,1,Int::from_data(Data::from([(item.labelasi64).elem()]),&self.device) }) .collect(); letimages=items .into_iter() .map(|item|{ letimage_data=item.image.iter().copied().flatten().collect::Vecf32(); letmutinput:TensorB,3= Tensor::from_floats(image_data.as_slice(),&self.device).reshape([1,28,28]); //Normalizetheinput input=((input/255)-0.1307)/0.3081; input }) .collect(); letimages=Tensor::stack(images,0); lettargets=Tensor::cat(targets,0); ClassificationBatch{images,targets} } } 3.2 迭代和损失策略我们这里使用经典的CrossEntropyLoss损失函数优化器使用AdaGradTrainStep是模型用于训练的trait，ValidStep是模型用于验证的trait，我们要为自动生成的Model，实现这两个trait，才能开始训练implB:BackendModel{ pubfnforward_classification( &self, images:TensorB,4, targets:TensorB,1,Int, )-ClassificationOutput{ letoutput=self.forward(images); letloss=CrossEntropyLossConfig::new() .init(&output.device()) .forward(output.clone(),targets.clone()); ClassificationOutput::new(loss,output,targets) } } implB:AutodiffBackendTrainStepClassificationBatch,ClassificationOutputforModel{ fnstep(&self,batch:ClassificationBatch)-TrainOutputClassificationOutput{ letitem=self.forward_classification(batch.images,batch.targets); TrainOutput::new(self,item.loss.backward(),item) } } implB:BackendValidStepClassificationBatch,ClassificationOutputforModel{ fnstep(&self,batch:ClassificationBatch)-ClassificationOutput{ self.forward_classification(batch.images,batch.targets) } } 3.3 训练过程，和添加监控下面是我们的整个训练过程，套路和torch一模一样我们将训练用的配置，和最终的模型，保存到文件中加入burn自带的监控，方便我们观察准确率和损失的变化。我们这里用new方法加载默认参数的模型，也就是没训练过的模型，而不是defaultpubfntrainB:AutodiffBackend(config:TrainingConfig,device:B::Device){ std::fs::create_dir_all(ARTIFACT_DIR).ok(); config .save(format!("{ARTIFACT_DIR}/config.json")) .expect("Configshouldbesavedsuccessfully"); B::seed(config.seed); letdataset_train=MNISTDataset::train(); letdataset_test=MNISTDataset::train(); //Dataloaders letbatcher_train=ClassificationBatcher::::new(device.clone()); letbatcher_valid=ClassificationBatcher::B::InnerBackend::new(device.clone()); letdataloader_train=DataLoaderBuilder::new(batcher_train) .batch_size(config.batch_size) .shuffle(config.seed) .num_workers(config.num_workers) .build(dataset_train); letdataloader_test=DataLoaderBuilder::new(batcher_valid) .batch_size(config.batch_size) .num_workers(config.num_workers) .build(dataset_test); //Learnerconfig letlearner=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 letnow=Instant::now(); letmodel_trained=learner.fit(dataloader_train,dataloader_test); letelapsed=now.elapsed().as_secs(); println!("Trainingcompletedin{}m{}s",(elapsed/60),elapsed%60); model_trained .save_file(format!("{ARTIFACT_DIR}/model"),&CompactRecorder::new()) .expect("Trainedmodelshouldbesavedsuccessfully"); } 4. 开始训练正式训练之前，需要配置一下训练相关的参数，将这个cfg传入train即可。 letmutcfg=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::AutodiffNdArray(cfg,NdArrayDevice::Cpu); run起程序， 风扇开始呼呼转动，窗口直接黑屏卡死，cpu监控拉满： image.png4.2 GPU接下来用我们的笔记本自带的垃圾卡跑一下（一般的a卡和i卡都可以跑） train::AutodiffWgpu(cfg,WgpuDevice::default()); 看一下监控的变化： 预计用时37分分钟准确率增加明显image.png再看一下损失变化情况，降低的也很好 看一下显卡的使用情况，a卡已经被利用起来了 最后看一下cpu的负载，有增加，但不多 image.png4.3 libtorch假如你很富有，有n卡，或者m1的mac，那么你同样可以用这种方式训练。 比如用m1芯片： train::AutodiffLibTorch(cfg,LibTorchDevice::Mps); 或者CUDA train::AutodiffLibTorch(cfg,LibTorchDevice::Cuda(0)); 我这里就不测试了 尾语 burn的能力不仅限于上面的种种，它允许自定义设备的接入。但是目前还处在较为初期的阶段，像我们pytorch例子中的nll_loss和Adadelta都是不支持的。 前几年挖矿，和最近的ai的尽头是电力和算力，都在大力拉升显卡的价格，国内显卡和好的芯片价格高得离谱，还要处处被卡脖子，哎~ 最后祝愿每个点赞收藏的帅哥美女都用上好卡。 阅读原文