OBJNULL/Dockerized-Rust-ML
Template
1
0
Fork 0
generated from OBJNULL/Dockerized-Rust
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Code Cleanup and Store in Out folder

Browse source
This commit is contained in:
Maddox Werts 2025-03-19 16:00:27 -04:00
parent 46149b885c
commit b9b07c4b75
1 changed files with 8 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
project/src/neural.rs
View file

@ -11,6 +11,9 @@ use burn::{
optim::AdamConfig,
};
// Constants
const MODEL_DIRECTORY: &str = "./out";
// Structures
pub struct NeuralNetwork {
mode: OperationMode,
@ -26,25 +29,23 @@ impl NeuralNetwork {
// Functions
fn train(&self) {
// Creating Backend
type MyBackend = Wgpu<f32, i32>;
type MyAutodiffBackend = Autodiff<MyBackend>;
// Create a default Wgpu device
let device = burn::backend::wgpu::WgpuDevice::default();
// All the training artifacts will be saved in this directory
let artifact_dir = "/tmp/guide";
// Train the model
training::train::<MyAutodiffBackend>(
artifact_dir,
MODEL_DIRECTORY,
training::TrainingConfig::new(model::ModelConfig::new(10, 512), AdamConfig::new()),
device.clone(),
);
// Infer the model
infrence::infer::<MyBackend>(
artifact_dir,
MODEL_DIRECTORY,
device,
burn::data::dataset::vision::MnistDataset::test()
.get(42)
@ -52,16 +53,13 @@ impl NeuralNetwork {
);
}
fn infer(&self) {
// Creating Backend
type MyBackend = Wgpu<f32, i32>;
let device = burn::backend::wgpu::WgpuDevice::default();
// All the training artifacts are saved in this directory
let artifact_dir = "/tmp/guide";
// Infer the model
infrence::infer::<MyBackend>(
artifact_dir,
MODEL_DIRECTORY,
device,
burn::data::dataset::vision::MnistDataset::test()
.get(42)