update modelwriter for rmsnorm

Tencent · Sep 3, 2024 · 1af1080 · 1af1080
1 parent 204583b
commit 1af1080
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;