add wrap_module in lib.rs

Author: 张扬发

src/lib.rs

@@ -3,6 +3,8 @@
 use pyo3::prelude::*;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::types::PyModule;
+use pyo3::wrap_pymodule;
+
 
 use candle_core::{DType, Device, Tensor};
 use candle_nn::{
@@ -11,11 +13,10 @@ use candle_nn::{
 
 pub mod base;   
 pub mod mlp;
-use mlp::MLP;
 
 
 #[pymodule]
 fn rust_mlp(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<MLP>()?;
+    m.add_wrapped(wrap_pymodule!(mlp::mlp))?;
     Ok(())
 }

src/mlp.rs

@@ -7,14 +7,14 @@ use candle_nn::{
 };
 use crate::base::Model;
 #[pyclass]
-pub struct MLP {
+pub struct PyMLP {
     model: SimpleNN,
     optimizer: candle_nn::AdamW,
     device: Device,
 }
 
 #[pymethods]
-impl MLP {
+impl PyMLP {
     #[new]
     fn new(input_dim: usize, lr: f64) -> PyResult<Self> {
         let device = Device::Cpu;
@@ -39,7 +39,7 @@ impl MLP {
         Ok(())
     }
 }
-impl Model for MLP {
+impl Model for PyMLP {
     fn train(&mut self, x: Vec<Vec<f32>>, y: Vec<f32>, epochs: usize) -> PyResult<()> {
         let n = x.len();
         let d = x[0].len();
@@ -101,4 +101,10 @@ impl Module for SimpleNN {
         let x = self.fc2.forward(&x)?;
         Ok(x)
     }
+}
+
+#[pymodule]
+pub fn mlp(m: &Bound<'_, PyModule>)-> PyResult<()>{
+    m.add_class::<PyMLP>()?;
+    Ok(())
 }

test.ipynb

@@ -46,45 +46,45 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 7,
    "id": "db84d87a",
    "metadata": {},
    "outputs": [],
    "source": [
     "input_dim = len(X_train[0])\n",
-    "model = rust_mlp.MLP(input_dim,lr = 0.01)"
+    "model = rust_mlp.mlp.PyMLP(input_dim,lr = 0.01)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 8,
    "id": "975ccf7f",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Loss: 11.602837\n",
-      "Loss: 9.104499\n",
-      "Loss: 7.028557\n",
-      "Loss: 5.354487\n",
-      "Loss: 4.0523024\n",
-      "Loss: 3.0819073\n",
-      "Loss: 2.3972578\n",
-      "Loss: 1.9500879\n",
-      "Loss: 1.6904843\n",
-      "Loss: 1.5687597\n",
-      "Loss: 1.5355564\n",
-      "Loss: 1.5495899\n",
-      "Loss: 1.5760833\n",
-      "Loss: 1.5907441\n",
-      "Loss: 1.5783068\n",
-      "Loss: 1.5354254\n",
-      "Loss: 1.4655299\n",
-      "Loss: 1.3767947\n",
-      "Loss: 1.2791892\n",
-      "Loss: 1.1817851\n"
+      "Loss: 6.421524\n",
+      "Loss: 4.9037886\n",
+      "Loss: 3.821808\n",
+      "Loss: 3.0536802\n",
+      "Loss: 2.4780674\n",
+      "Loss: 2.0242903\n",
+      "Loss: 1.6653357\n",
+      "Loss: 1.3944621\n",
+      "Loss: 1.2099363\n",
+      "Loss: 1.1070489\n",
+      "Loss: 1.074543\n",
+      "Loss: 1.0935023\n",
+      "Loss: 1.1390615\n",
+      "Loss: 1.185772\n",
+      "Loss: 1.2136548\n",
+      "Loss: 1.2125317\n",
+      "Loss: 1.1814777\n",
+      "Loss: 1.1264257\n",
+      "Loss: 1.0574386\n",
+      "Loss: 0.9845169\n"
      ]
     }
    ],
@@ -94,15 +94,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 9,
    "id": "fd3e1b4c",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Evaluation Loss: 1.2425249\n"
+      "Evaluation Loss: 0.9266705\n"
      ]
     }
    ],