增加返回权重的方法

zjhellofss · Feb 10, 2024 · bc6605e · bc6605e
1 parent e6a99ac
commit bc6605e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/include/layer/abstract/param_layer.hpp b/include/layer/abstract/param_layer.hpp
@@ -122,6 +122,8 @@ class ParamLayer : public Layer<float> {
    */
   void set_bias(const std::vector<std::shared_ptr<Tensor<float>>>& bias) override;
 
+  std::shared_ptr<Tensor<float>> weight(int32_t index) const;
+
  protected:
   std::vector<std::shared_ptr<Tensor<float>>> weights_;
   std::vector<std::shared_ptr<Tensor<float>>> bias_;

diff --git a/source/layer/abstract/param_layer.cpp b/source/layer/abstract/param_layer.cpp
@@ -94,6 +94,11 @@ void ParamLayer::set_weights(const std::vector<float>& weights) {
   }
 }
 
+std::shared_ptr<Tensor<float>> ParamLayer::weight(int32_t index) const {
+  CHECK_LE(index, this->weights_.size());
+  return this->weights_.at(index);
+}
+
 void ParamLayer::set_bias(const std::vector<float>& bias) {
   size_t bias_size = 0;
   const size_t elem_size = bias.size();