add class

Author: 张扬发

src/lib.rs

@@ -10,12 +10,14 @@ use candle_nn::{
 };
 mod mytest;
 use mytest::Number;
-
+mod mlp;
+use mlp::MLP;
 #[pymodule]
 fn rust_mlp(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
     m.add_function(wrap_pyfunction!(let_me_try, m)?)?;
     m.add_class::<Number>()?;
+    m.add_class::<MLP>()?;
     Ok(())
 }
 

src/mlp.rs

@@ -0,0 +1,87 @@
+use pyo3::prelude::*;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::types::PyModule;
+use candle_core::{DType, Device, Tensor};
+use candle_nn::{
+    linear, Linear, Module, VarBuilder, VarMap, Optimizer,
+};
+
+#[pyclass]
+pub struct MLP {
+    model: SimpleNN,
+    optimizer: candle_nn::AdamW,
+    device: Device,
+}
+
+#[pymethods]
+impl MLP {
+    #[new]
+    fn new(input_dim: usize, lr: f64) -> PyResult<Self> {
+        let device = Device::Cpu;
+        let varmap = VarMap::new();
+        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+        let model = SimpleNN::new(input_dim, vb)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        let optimizer = candle_nn::AdamW::new_lr(varmap.all_vars(), lr)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        Ok(Self { model, optimizer, device })
+    }
+
+    fn train(&mut self, x: Vec<Vec<f32>>, y: Vec<f32>, epochs: usize) -> PyResult<()> {
+        let n = x.len();
+        let d = x[0].len();
+        let x_flat = x.into_iter().flatten().collect::<Vec<_>>();
+        let x_tensor = Tensor::from_vec(x_flat, (n, d), &self.device)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        let y_tensor = Tensor::from_vec(y, (n, 1), &self.device)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        for _ in 0..epochs {
+            let output = self.model.forward(&x_tensor)
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            let loss = candle_nn::loss::mse(&output, &y_tensor)
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            self.optimizer.backward_step(&loss)
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        }
+        Ok(())
+    }
+
+    fn evaluate(&self, x: Vec<Vec<f32>>, y: Vec<f32>) -> PyResult<f32> {
+        let n = x.len();
+        let d = x[0].len();
+        let x_flat = x.into_iter().flatten().collect::<Vec<_>>();
+        let x_tensor = Tensor::from_vec(x_flat, (n, d), &self.device)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        let y_tensor = Tensor::from_vec(y, (n, 1), &self.device)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        let output = self.model.forward(&x_tensor)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        let loss = candle_nn::loss::mse(&output, &y_tensor)
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+        Ok(loss.to_scalar::<f32>()
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?)
+    }
+}
+
+#[derive(Debug)]
+struct SimpleNN {
+    fc1: Linear,
+    fc2: Linear,
+}
+
+impl SimpleNN {
+    fn new(in_dim: usize, vb: VarBuilder) -> candle_core::Result<Self> {
+        let fc1 = linear(in_dim, 64, vb.pp("fc1"))?;
+        let fc2 = linear(64, 1, vb.pp("fc2"))?;
+        Ok(Self { fc1, fc2 })
+    }
+}
+
+impl Module for SimpleNN {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let x = self.fc1.forward(xs)?;
+        let x = x.relu()?;
+        let x = self.fc2.forward(&x)?;
+        Ok(x)
+    }
+}

src/mytest.rs

@@ -1,4 +1,5 @@
 use pyo3::prelude::*;
+use pyo3::types::PyModule;
 
 #[pyclass]
 pub struct Number(i32);
@@ -9,4 +10,5 @@ impl Number {
     fn new(value: i32) -> Self {
         Number(value)
     }
-}
+}
+

test.ipynb

@@ -12,12 +12,45 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "bedc6943",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<Number at 0x106bce4f0>"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "x = rust_mlp.Number(42)\n",
+    "x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "db84d87a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<MLP at 0x106bdd740>"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "x = rust_mlp.Number(42)"
+    "rust_mlp.MLP(64,0.01)"
    ]
   }
  ],