[NeoML] fix CTransformerEncoderLayer reconstruction in change preNorm (…

…#1055) Signed-off-by: Kirill Golikov <kirill.golikov@abbyy.com>
neoml-lib · Apr 26, 2024 · 6ccc1da · 6ccc1da
1 parent fafb496
commit 6ccc1da
Showing 1 changed file with 24 additions and 2 deletions.
diff --git a/NeoML/src/Dnn/Layers/TransformerLayer.cpp b/NeoML/src/Dnn/Layers/TransformerLayer.cpp
@@ -95,8 +95,30 @@ void CTransformerEncoderLayer::Serialize( CArchive& archive )
 
 void CTransformerEncoderLayer::SetPreNorm( bool _preNorm )
 {
-	preNorm = _preNorm;
-	DeleteAllLayers();
+	if( preNorm != _preNorm ) {
+		preNorm = _preNorm;
+
+		const auto headCount = GetHeadCount();
+		const auto hiddenSize = GetHiddenSize();
+		const auto dropoutRate = GetDropoutRate();
+		const auto attentionDropoutRate = GetSelfAttentionDropoutRate();
+		const auto feedForwardSize = GetFeedForwardSize();
+		const auto maskType = GetMaskType();
+
+		const auto& activationLayer = dynamic_cast<const IActivationLayer&>( *GetLayer( activationName ) );
+		const auto activation = activationLayer.GetDesc();
+
+		DeleteAllLayers();
+		buildLayer();
+
+		SetHeadCount( headCount );
+		SetHiddenSize( hiddenSize );
+		SetDropoutRate( dropoutRate );
+		SetSelfAttentionDropoutRate( attentionDropoutRate );
+		SetFeedForwardSize( feedForwardSize );
+		SetMaskType( maskType );
+		SetActivation( activation );
+	}
 }
 
 void CTransformerEncoderLayer::SetHeadCount( int headCount )