add class

Author: 张扬发

README.md

@@ -13,29 +13,40 @@ pip install Path/to/*.whl
 - 在python环境中使用
 ```Python
 import rust_mlp
-rust_mlp.let_me_try()
+import json
+with open('fetch_california_housing.json', encoding='utf-8') as f:
+    data = json.load(f) 
+X_train = data['X_train']
+X_test  = data['X_test']
+y_train = data['y_train']
+y_test  = data['y_test']
+
+input_dim = len(X_train[0])
+model = rust_mlp.MLP(input_dim,lr = 0.01)
+model.train(X_train,y_train, epochs = 20)
+model.evaluate(X_test,y_test)
 ```
 预期结果如下：
 ```
-Epoch 10 | Train Loss: 1.628042
-Epoch 20 | Train Loss: 1.206229
-Epoch 30 | Train Loss: 0.876774
-Epoch 40 | Train Loss: 0.737402
-Epoch 50 | Train Loss: 0.644500
-Epoch 60 | Train Loss: 0.576381
-Epoch 70 | Train Loss: 0.524079
-Epoch 80 | Train Loss: 0.482977
-Epoch 90 | Train Loss: 0.451547
-Epoch 100 | Train Loss: 0.431264
-Epoch 110 | Train Loss: 0.417604
-Epoch 120 | Train Loss: 0.407959
-Epoch 130 | Train Loss: 0.400804
-Epoch 140 | Train Loss: 0.395103
-Epoch 150 | Train Loss: 0.390289
-Epoch 160 | Train Loss: 0.386080
-Epoch 170 | Train Loss: 0.382233
-Epoch 180 | Train Loss: 0.378714
-Epoch 190 | Train Loss: 0.375436
-Epoch 200 | Train Loss: 0.372300
-Test Loss: 0.385307
+Loss: 4.4525466
+Loss: 3.3261764
+Loss: 2.539338
+Loss: 2.008424
+Loss: 1.6563666
+Loss: 1.4215927
+Loss: 1.2626717
+Loss: 1.1582772
+Loss: 1.0978976
+Loss: 1.0723116
+Loss: 1.0687258
+Loss: 1.073826
+Loss: 1.0768875
+Loss: 1.0705063
+Loss: 1.0518045
+Loss: 1.0223864
+Loss: 0.9862392
+Loss: 0.9472292
+Loss: 0.9076066
+Loss: 0.86793756
+Evaluation Loss: 0.82283765
 ```

src/lib.rs

@@ -8,128 +8,10 @@ use candle_core::{DType, Device, Tensor};
 use candle_nn::{
     linear, Linear, Module, VarBuilder, VarMap, Optimizer,
 };
-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(())
-}
-
-#[pyfunction]
-fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
-    Ok((a + b).to_string())
-}
-
-#[pyfunction]
-fn let_me_try() -> PyResult<()> {
-    run_training()
-        .map_err(|e| PyRuntimeError::new_err(e.to_string()))
-}
-
-fn run_training() -> anyhow::Result<()> {
-    let device = Device::Cpu;
-
-    let file = std::fs::File::open("fetch_california_housing.json")?;
-    let reader = std::io::BufReader::new(file);
-    let data: Data = serde_json::from_reader(reader)?;
-
-    let train_d1 = data.X_train.len();
-    let train_d2 = data.X_train[0].len();
-    let test_d1 = data.X_test.len();
-    let test_d2 = data.X_test[0].len();
-
-    let x_train_vec = data.X_train.into_iter().flatten().collect::<Vec<_>>();
-    let x_test_vec = data.X_test.into_iter().flatten().collect::<Vec<_>>();
-
-    let y_train_vec = data.y_train;
-    let y_test_vec = data.y_test;
-
-    let x_train = Tensor::from_vec(x_train_vec, (train_d1, train_d2), &device)?;
-    let y_train = Tensor::from_vec(y_train_vec, (train_d1, 1), &device)?;
-    let x_test = Tensor::from_vec(x_test_vec, (test_d1, test_d2), &device)?;
-    let y_test = Tensor::from_vec(y_test_vec, (test_d1, 1), &device)?;
-
-    let varmap = VarMap::new();
-    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
-
-    let model = SimpleNN::new(train_d2, vb)?;
-
-    let mut optimizer =
-        candle_nn::AdamW::new_lr(varmap.all_vars(), 1e-2)?;
-
-    train_model(&model, &x_train, &y_train, &mut optimizer, 200)?;
-    evaluate_model(&model, &x_test, &y_test)?;
-
-    Ok(())
-}
-
-#[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)
-    }
-}
-
-fn train_model(
-    model: &SimpleNN,
-    x_train: &Tensor,
-    y_train: &Tensor,
-    optimizer: &mut candle_nn::AdamW,
-    epochs: usize,
-) -> anyhow::Result<()> {
-    for epoch in 0..epochs {
-        let output = model.forward(x_train)?;
-        let loss = candle_nn::loss::mse(&output, y_train)?;
-        optimizer.backward_step(&loss)?;
-
-        if (epoch + 1) % 10 == 0 {
-            println!(
-                "Epoch {} | Train Loss: {:.6}",
-                epoch + 1,
-                loss.to_scalar::<f32>()?
-            );
-        }
-    }
-    Ok(())
-}
-
-fn evaluate_model(
-    model: &SimpleNN,
-    x_test: &Tensor,
-    y_test: &Tensor,
-) -> anyhow::Result<()> {
-    let output = model.forward(x_test)?;
-    let loss = candle_nn::loss::mse(&output, y_test)?;
-    println!("Test Loss: {:.6}", loss.to_scalar::<f32>()?);
-    Ok(())
-}
-
-#[derive(Debug, serde::Deserialize)]
-struct Data {
-    X_train: Vec<Vec<f32>>,
-    X_test: Vec<Vec<f32>>,
-    y_train: Vec<f32>,
-    y_test: Vec<f32>,
-}
+}

src/mlp.rs

@@ -40,13 +40,15 @@ impl MLP {
                 .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()))?;
+            println!("Loss: {}", loss.to_scalar::<f32>()
+                .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> {
+    fn evaluate(&self, x: Vec<Vec<f32>>, y: Vec<f32>) -> PyResult<()> {
         let n = x.len();
         let d = x[0].len();
         let x_flat = x.into_iter().flatten().collect::<Vec<_>>();
@@ -58,8 +60,9 @@ impl MLP {
             .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()))?)
+        println!("Evaluation Loss: {}", loss.to_scalar::<f32>()
+            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?);
+        Ok(())
     }
 }
 

test.ipynb

@@ -10,47 +10,104 @@
     "import rust_mlp"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "4ca52908",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "bedc6943",
+   "id": "34d74ef8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open('fetch_california_housing.json', encoding='utf-8') as f:\n",
+    "    data = json.load(f) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "d0bb55a4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "X_train = data['X_train']\n",
+    "X_test  = data['X_test']\n",
+    "y_train = data['y_train']\n",
+    "y_test  = data['y_test']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "db84d87a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_dim = len(X_train[0])\n",
+    "model = rust_mlp.MLP(input_dim,lr = 0.01)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "975ccf7f",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "<Number at 0x106bce4f0>"
-      ]
-     },
-     "execution_count": 3,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loss: 4.4525466\n",
+      "Loss: 3.3261764\n",
+      "Loss: 2.539338\n",
+      "Loss: 2.008424\n",
+      "Loss: 1.6563666\n",
+      "Loss: 1.4215927\n",
+      "Loss: 1.2626717\n",
+      "Loss: 1.1582772\n",
+      "Loss: 1.0978976\n",
+      "Loss: 1.0723116\n",
+      "Loss: 1.0687258\n",
+      "Loss: 1.073826\n",
+      "Loss: 1.0768875\n",
+      "Loss: 1.0705063\n",
+      "Loss: 1.0518045\n",
+      "Loss: 1.0223864\n",
+      "Loss: 0.9862392\n",
+      "Loss: 0.9472292\n",
+      "Loss: 0.9076066\n",
+      "Loss: 0.86793756\n"
+     ]
     }
    ],
    "source": [
-    "x = rust_mlp.Number(42)\n",
-    "x"
+    "model.train(X_train,y_train, epochs = 20)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "id": "db84d87a",
+   "execution_count": 7,
+   "id": "fd3e1b4c",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "<MLP at 0x106bdd740>"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Evaluation Loss: 0.82283765\n"
+     ]
     }
    ],
    "source": [
-    "rust_mlp.MLP(64,0.01)"
+    "model.evaluate(X_test,y_test)"
    ]
   }
  ],