Use new methods to remove indirection

EricLBuehler · Sep 11, 2023 · 55a5bbe · 55a5bbe
1 parent 4fe27dd
commit 55a5bbe
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 54 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -106,28 +106,15 @@ pub trait Conv1dLayerLike: Module {
     fn config(&self) -> &Conv1dConfig;
 }
 
-#[derive(Debug)]
-pub struct Conv1dWithWB {
-    pub layer: Conv1d,
-    pub weights: Tensor,
-    pub bias: Option<Tensor>,
-}
-
-impl Module for Conv1dWithWB {
-    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
-        self.layer.forward(xs)
-    }
-}
-
-impl Conv1dLayerLike for Conv1dWithWB {
+impl Conv1dLayerLike for Conv1d {
     fn config(&self) -> &Conv1dConfig {
-        self.layer.config()
+        self.config()
     }
     fn weight(&self) -> &Tensor {
-        &self.weights
+        self.weight()
     }
     fn bias(&self) -> Option<&Tensor> {
-        self.bias.as_ref()
+        self.bias()
     }
 }
 
@@ -138,28 +125,15 @@ pub trait Conv2dLayerLike: Module {
     fn config(&self) -> &Conv2dConfig;
 }
 
-#[derive(Debug)]
-pub struct Conv2dWithWB {
-    pub layer: Conv2d,
-    pub weights: Tensor,
-    pub bias: Option<Tensor>,
-}
-
-impl Module for Conv2dWithWB {
-    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
-        self.layer.forward(xs)
-    }
-}
-
-impl Conv2dLayerLike for Conv2dWithWB {
+impl Conv2dLayerLike for Conv2d {
     fn config(&self) -> &Conv2dConfig {
-        self.layer.config()
+        self.config()
     }
     fn weight(&self) -> &Tensor {
-        &self.weights
+        self.weight()
     }
     fn bias(&self) -> Option<&Tensor> {
-        self.bias.as_ref()
+        self.bias()
     }
 }
 

diff --git a/tests/conv1d.rs b/tests/conv1d.rs
@@ -4,8 +4,7 @@ fn conv1d() -> candle_core::Result<()> {
 
     use candle_core::{DType, Device, Result, Tensor};
     use candle_lora::{
-        loraconv1d::LoraConv1dConfig, Conv1dLayerLike, Conv1dWithWB, Lora, NewLayers,
-        SelectedLayers,
+        loraconv1d::LoraConv1dConfig, Conv1dLayerLike, Lora, NewLayers, SelectedLayers,
     };
     use candle_nn::{init, Conv1d, Conv1dConfig, Module, VarMap};
 
@@ -55,18 +54,12 @@ fn conv1d() -> candle_core::Result<()> {
         &device,
     )?;
 
-    let conv = Conv1dWithWB {
-        layer: Conv1d::new(
+    let mut model = Model {
+        conv: Box::new(Conv1d::new(
             conv_weight.clone(),
             Some(conv_bias.clone()),
             Conv1dConfig::default(),
-        ),
-        weights: conv_weight,
-        bias: Some(conv_bias),
-    };
-
-    let mut model = Model {
-        conv: Box::new(conv),
+        )),
     };
 
     let dummy_image = Tensor::zeros((1, 10, 10), DType::F32, &device)?;

diff --git a/tests/conv2d.rs b/tests/conv2d.rs
@@ -4,8 +4,7 @@ fn conv2d() -> candle_core::Result<()> {
 
     use candle_core::{DType, Device, Result, Tensor};
     use candle_lora::{
-        loraconv2d::LoraConv2dConfig, Conv2dLayerLike, Conv2dWithWB, Lora, NewLayers,
-        SelectedLayers,
+        loraconv2d::LoraConv2dConfig, Conv2dLayerLike, Lora, NewLayers, SelectedLayers,
     };
     use candle_nn::{init, Conv2d, Conv2dConfig, Module, VarMap};
 
@@ -66,14 +65,12 @@ fn conv2d() -> candle_core::Result<()> {
         &device,
     )?;
 
-    let conv = Conv2dWithWB {
-        layer: Conv2d::new(conv_weight.clone(), Some(conv_bias.clone()), cfg),
-        weights: conv_weight,
-        bias: Some(conv_bias),
-    };
-
     let mut model = Model {
-        conv: Box::new(conv),
+        conv: Box::new(Conv2d::new(
+            conv_weight.clone(),
+            Some(conv_bias.clone()),
+            cfg,
+        )),
     };
 
     let shape = [2, in_channels, 20, 20]; //(BS, K, X, Y)