update modelwriter for rmsnorm (#5676)

Tencent · Sep 4, 2024 · 21e54d8 · 21e54d8
1 parent 204583b
commit 21e54d8
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/tools/modelwriter.h b/tools/modelwriter.h
@@ -99,6 +99,7 @@
 #include "layer/reorg.h"
 #include "layer/requantize.h"
 #include "layer/reshape.h"
+#include "layer/rmsnorm.h"
 #include "layer/rnn.h"
 #include "layer/roialign.h"
 #include "layer/roipooling.h"
@@ -2313,6 +2314,17 @@ int ModelWriter::save(const char* parampath, const char* binpath)
  fprintf_param_value(" 2=%d", c)
  fprintf_param_value(" 3=%d", permute)
  }
+ else if (layer->type == "RMSNorm")
+ {
+ ncnn::RMSNorm* op = (ncnn::RMSNorm*)layer;
+ ncnn::RMSNorm* op_default = (ncnn::RMSNorm*)layer_default;
+
+ fprintf_param_value(" 0=%d", affine_size)
+ fprintf_param_value(" 1=%e", eps)
+ fprintf_param_value(" 2=%d", affine)
+
+ fwrite_weight_data(op->gamma_data, bp);
+ }
  else if (layer->type == "RNN")
  {
  ncnn::RNN* op = (ncnn::RNN*)layer;