[NeoML] CTiedEmbeddingsLayer is learnable with empty paramBlobs (#1106)

Signed-off-by: Kirill Golikov <kirill.golikov@abbyy.com>

Author: favorart

NeoML/include/NeoML/Dnn/Dnn.h

@@ -244,8 +244,10 @@ class NEOML_API CBaseLayer : public virtual IObject {
 	bool IsBackwardPerformed() const;
 	// Indicates that backpropagation must be performed for the layer when Learn method is called
 	bool IsBackwardNeeded() const;
-	// Layer may contain empty paramBlob of given index
-	virtual bool ContainsEmptyParamBlob( int ) const { return false; }
+	// Layer may contain null paramBlob of given index, specialization for transferParamsBlob
+	virtual bool ContainsNullParamBlob( int ) const { return false; }
+	// Special case, specialization for transferParamsBlob
+	virtual bool IsLearnableWithEmptyParamBlobs() const { return false; }
 	// Gets a pointer to the layer connected to the given input
 	CBaseLayer* GetInputLayer(int input) { return inputLinks[input].Layer; }
 	const CBaseLayer* GetInputLayer(int input) const { return inputLinks[input].Layer; }

NeoML/include/NeoML/Dnn/DnnSolver.h

@@ -29,7 +29,7 @@ class NEOML_API CDnnSolver : virtual public IObject {
 public:
 	// Stores the calculated values of layer parameters gradients for further use in Train method
 	// forSharedWeightsLayer=true should only be used within layers that share weights with other layers.
-	void AddDiff( CBaseLayer* layer, const CObjectArray<CDnnBlob>& paramDiffBlobs, 
+	void AddDiff( const CBaseLayer* layer, const CObjectArray<CDnnBlob>& paramDiffBlobs, 
 		bool sharedWeights = false );
 
 	// Modifies the trainable parameters of the network layers, 
@@ -98,8 +98,8 @@ class NEOML_API CDnnSolver : virtual public IObject {
 	// Used in the inheriting classes
 	CMap<CString, CObjectArray<CDnnBlob>> layerToGradientHistory;
 	// Layers which require reduction across distributed solver
-	CHashTable<CBaseLayer*> layersToReduce; // Fast check if layer is included already
-	CArray<CBaseLayer*> reduceOrder; // Correct order across all of the distributed nets
+	CHashTable<const CBaseLayer*> layersToReduce; // Fast check if layer is included already
+	CArray<const CBaseLayer*> reduceOrder; // Correct order across all of the distributed nets
 
 	// Averages weights over all threads
 	void allReduce( float distributedCoeff );

NeoML/include/NeoML/Dnn/Layers/ChannelwiseWith1x1Layer.h

@@ -58,7 +58,7 @@ class NEOML_API CChannelwiseWith1x1Layer : public CBaseLayer {
 	void RunOnce() override;
 	void BackwardOnce() override { NeoAssert( false ); }
 	// Specialization for transferParamsBlob
-	bool ContainsEmptyParamBlob( int i ) const override
+	bool ContainsNullParamBlob( int i ) const override
 		{ return paramBlobs[i] == nullptr && ( i == P_ChannelwiseFreeTerm || i == P_ConvFreeTerm ); }
 
 private:

NeoML/include/NeoML/Dnn/Layers/MobileNetV2BlockLayer.h

@@ -75,7 +75,7 @@ class NEOML_API CMobileNetV2BlockLayer : public CBaseLayer {
 	void RunOnce() override;
 	void BackwardOnce() override { NeoAssert( false ); }
 	// Specialization for transferParamsBlob
-	bool ContainsEmptyParamBlob( int i ) const override
+	bool ContainsNullParamBlob( int i ) const override
 		{ return !paramBlobs[i] && ( i == P_ChannelwiseFreeTerm || i == P_DownFreeTerm || i == P_ExpandFreeTerm ); }
 
 private:

NeoML/include/NeoML/Dnn/Layers/MobileNetV3BlockLayer.h

@@ -59,7 +59,7 @@ class NEOML_API CMobileNetV3PreSEBlockLayer : public CBaseLayer {
 	void RunOnce() override;
 	void BackwardOnce() override { NeoAssert( false ); }
 	// Specialization for transferParamsBlob
-	bool ContainsEmptyParamBlob( int i ) const override
+	bool ContainsNullParamBlob( int i ) const override
 		{ return paramBlobs[i] == nullptr && ( i == P_ChannelwiseFreeTerm || i == P_ExpandFreeTerm ); }
 
 private:
@@ -115,7 +115,7 @@ class NEOML_API CMobileNetV3PostSEBlockLayer : public CBaseLayer {
 	void RunOnce() override;
 	void BackwardOnce() override { NeoAssert( false ); }
 	// Specialization for transferParamsBlob
-	bool ContainsEmptyParamBlob( int i ) const override
+	bool ContainsNullParamBlob( int i ) const override
 		{ return paramBlobs[i] == nullptr && ( i == P_DownFreeTerm ); }
 
 private:

NeoML/include/NeoML/Dnn/Layers/TiedEmbeddingsLayer.h

@@ -27,7 +27,8 @@ class CMultichannelLookupLayer;
 class NEOML_API CTiedEmbeddingsLayer : public CBaseLayer {
 	NEOML_DNN_LAYER( CTiedEmbeddingsLayer )
 public:
-	explicit CTiedEmbeddingsLayer( IMathEngine& mathEngine );
+	explicit CTiedEmbeddingsLayer( IMathEngine& mathEngine ) :
+		CBaseLayer( mathEngine, "CTiedEmbeddingsLayer", /*isLearnable*/true ) {}
 
 	void Serialize( CArchive& archive ) override;
 
@@ -54,17 +55,19 @@ class NEOML_API CTiedEmbeddingsLayer : public CBaseLayer {
 	void LearnOnce() override;
 	int BlobsForBackward() const override { return 0; }
 	int BlobsForLearn() const override { return TInputBlobs; }
+	// Special case, specialization for transferParamsBlob
+	bool IsLearnableWithEmptyParamBlobs() const override { return true; }
 
 private:
 	// Path for embedding layer from which matrix is taken
 	// Now it contains the path as array
 	// So in case of no composite layer it is gonna be { "embeddingName" }
 	CArray<CString> embeddingPath;
 	// Channel index in embedding layer.
-	int channelIndex;
+	int channelIndex = 0;
 
 	const CDnnBlob* getEmbeddingsTable() const;
-	CMultichannelLookupLayer* getLookUpLayer() const;
+	const CMultichannelLookupLayer* getLookUpLayer() const;
 };
 
 // Tied embeddings.

NeoML/src/Dnn/BaseLayer.cpp

@@ -314,13 +314,17 @@ void CBaseLayer::transferParamsBlob( CBaseLayer& dist ) const
 		}
 	} else {
 		NeoAssertMsg( dist.paramBlobs.Size() == paramBlobs.Size(), "transferParamsBlob: It isn't a copy of the layer" );
+		if( IsLearnableWithEmptyParamBlobs() ) { // Special case is CTiedEmbeddingsLayer
+			NeoAssert( dist.IsLearnable() && paramBlobs.Size() == 0 );
+			return;
+		}
 
 		NeoAssertMsg( !dist.IsLearnable() || paramBlobs.Size() > 0,
 			"transferParamsBlob: The origin dnn should be trained and reshaped to create a reference dnn" );
 		// Create reference copy of dist.paramBlobs with shared buffer
 		// Takes a pointer to parent's blob to access memory
 		for( int j = 0; j < dist.paramBlobs.Size(); ++j ) {
-			if( ContainsEmptyParamBlob( j ) ) {
+			if( ContainsNullParamBlob( j ) ) {
 				dist.paramBlobs[j] = nullptr; // may contain empty parameter
 				continue;
 			}

NeoML/src/Dnn/DnnSolver.cpp

@@ -140,7 +140,7 @@ CDnnSolver::CDnnSolver( IMathEngine& _mathEngine ) :
 }
 
 // Calculates the layer parameter gradients to then use them in Train method
-void CDnnSolver::AddDiff( CBaseLayer* layer, const CObjectArray<CDnnBlob>& paramDiffBlobs, bool sharedWeights )
+void CDnnSolver::AddDiff( const CBaseLayer* layer, const CObjectArray<CDnnBlob>& paramDiffBlobs, bool sharedWeights )
 {
 	NeoAssert( layer != nullptr );
 	if( MathEngine().IsDistributed() && !layersToReduce.Has( layer ) ) {

NeoML/src/Dnn/Layers/TiedEmbeddingsLayer.cpp

@@ -21,12 +21,6 @@ limitations under the License.
 
 namespace NeoML {
 
-CTiedEmbeddingsLayer::CTiedEmbeddingsLayer( IMathEngine& mathEngine ) :
-	CBaseLayer( mathEngine, "CTiedEmbeddingsLayer", true ),
-	channelIndex( 0 )
-{
-}
-
 void CTiedEmbeddingsLayer::SetChannelIndex( int val )
 {
 	NeoAssert( val >= 0 );
@@ -131,7 +125,7 @@ void CTiedEmbeddingsLayer::LearnOnce()
 		diffBlob->Clear();
 	}
 
-	CMultichannelLookupLayer* embeddingsLayer = getLookUpLayer();
+	const CMultichannelLookupLayer* embeddingsLayer = getLookUpLayer();
 	CObjectArray<CDnnBlob> totalDiffBlobs;
 	const int channelsCount = embeddingsLayer->GetDimensions().Size();
 	for( int i = 0; i < channelsCount; i++ ) {
@@ -144,7 +138,7 @@ void CTiedEmbeddingsLayer::LearnOnce()
 		}
 	}
 
-	GetDnn()->GetSolver()->AddDiff( embeddingsLayer, totalDiffBlobs, true );
+	GetDnn()->GetSolver()->AddDiff( embeddingsLayer, totalDiffBlobs, /*sharedWeights*/true );
 }
 
 // Embeddings matrix
@@ -155,11 +149,10 @@ const CDnnBlob* CTiedEmbeddingsLayer::getEmbeddingsTable() const
 	return getLookUpLayer()->GetEmbeddings( channelIndex );
 }
 
-CMultichannelLookupLayer* CTiedEmbeddingsLayer::getLookUpLayer() const
+const CMultichannelLookupLayer* CTiedEmbeddingsLayer::getLookUpLayer() const
 {
-	CMultichannelLookupLayer* embeddingsLayer;
-	embeddingsLayer = CheckCast<CMultichannelLookupLayer>(
-		const_cast<CDnn*>(GetDnn())->GetLayer(embeddingPath).Ptr());
+	const CMultichannelLookupLayer* embeddingsLayer
+		= CheckCast<CMultichannelLookupLayer>( GetDnn()->GetLayer( embeddingPath ).Ptr() );
 	return embeddingsLayer;
 }