[NeoML] IShuffledBatchGenerator in CProblemSourceLayer (neoml-lib#1104)

Signed-off-by: Kirill Golikov <[email protected]>

Author: favorart

NeoML/include/NeoML/Dnn/Layers/FullyConnectedSourceLayer.h

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -78,7 +78,9 @@ class NEOML_API CFullyConnectedSourceLayer : public CFullyConnectedLayer {
 	bool isBatchLoaded( int index ) const;
 };
 
-NEOML_API CLayerWrapper<CFullyConnectedSourceLayer> FullyConnectedSource(
-	TBlobType labelType, int batchSize, int maxBatchCount, IProblem* problem );
+// Creates CFullyConnectedSourceLayer with the name
+NEOML_API CFullyConnectedSourceLayer* FullyConnectedSource( CDnn& dnn, const char* name,
+	TBlobType labelType, int batchSize, int maxBatchCount, IProblem* problem,
+	int numberOfElements = 1, bool isZeroFreeTerm = false );
 
 } // namespace NeoML

NeoML/include/NeoML/Dnn/Layers/ModelWrapperLayer.h

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -27,31 +27,46 @@ class CSourceLayer;
 class CSinkLayer;
 class CDnnTrainingModelWrapper;
 
+struct IShuffledBatchGenerator {
+	virtual ~IShuffledBatchGenerator() = default;
+
+	virtual const CArray<int>& GenerateBatchIndexes( int batchSize, bool batchShuffled ) = 0;
+	virtual bool HasUnseenElements() const = 0;
+	virtual void DeleteUnseenElement( int index ) = 0;
+};
+
+//---------------------------------------------------------------------------------------------------------------------
+
 // CProblemSourceLayer is a wrapper over the IProblem interface. 
 // On each iteration, it passes BatchSize vectors into the network for processing.
 class NEOML_API CProblemSourceLayer : public CBaseLayer {
 	NEOML_DNN_LAYER( CProblemSourceLayer )
 public:
-	explicit CProblemSourceLayer( IMathEngine& mathEngine );
+	explicit CProblemSourceLayer( IMathEngine& mathEngine ) :
+		CBaseLayer( mathEngine, "CCnnProblemSourceLayer", /*isLearnable*/false ) {}
 
 	void Serialize( CArchive& archive ) override;
 
 	int GetBatchSize() const { return batchSize; }
-	void SetBatchSize(int _batchSize);
+	void SetBatchSize( int batchSize );
 
 	// The filler for empty values that are not present in a sparse vector
 	float GetEmptyFill() const { return emptyFill; }
-	void SetEmptyFill(float _emptyFill) { NeoAssert(GetDnn() == 0); emptyFill = _emptyFill; }
+	void SetEmptyFill( float _emptyFill ) { NeoAssert( GetDnn() == nullptr ); emptyFill = _emptyFill; }
 
 	// You may only change the problem for the layer that is connected to a network
 	// if the number of classes and the number of input vectors stay the same
 	CPtr<const IProblem> GetProblem() const { return problem; }
-	void SetProblem(const CPtr<const IProblem>& _problem);
+	void SetProblem( const CPtr<const IProblem>& problem, bool shuffle = false, unsigned seed = 42 );
 
 	// Retrieves and sets the data type for class labels
 	TBlobType GetLabelType() const { return labelType; }
 	void SetLabelType( TBlobType newLabelType );
 
+	// Still not the end of an epoch
+	bool HasUnseenElements() const
+		{ return ( shuffled && shuffled->HasUnseenElements() ) || nextProblemIndex < problem->GetVectorCount(); }
+
 protected:
 	~CProblemSourceLayer() override = default;
 
@@ -60,34 +75,43 @@ class NEOML_API CProblemSourceLayer : public CBaseLayer {
 	void BackwardOnce() override;
 
 private:
-	float emptyFill;		// the empty values filler (for values not represented in a sparse vector)
-	int batchSize;			// the size of the batch passed to the network
-	int nextProblemIndex;	// the index of the next element in the problem to be passed
-	CPtr<const IProblem> problem;	// the classification problem the network is solving
-	TBlobType labelType;		// the data type for labels
-	CArray<float> exchangeBufs[3];
+	float emptyFill = 0; // the empty values filler (for values not represented in a sparse vector)
+	int batchSize = 1; // the size of the batch passed to the network
+	int nextProblemIndex = NotFound; // the index of the next element in the problem to be passed
+	TBlobType labelType = CT_Float; // the data type for labels
+	CPtr<const IProblem> problem; // the classification problem the network is solving
+	CPtrOwner<IShuffledBatchGenerator> shuffled; // if a shuffled batch input
+
+	enum { EB_Data, EB_Label, EB_Weight, EB_Count_ };
+	CArray<float> exchangeBufs[EB_Count_]{};
+
+	void fillExchangeBuffers( int shift, int index );
 };
 
-///////////////////////////////////////////////////////////////////////////////////////////////////////
+// Creates CProblemSourceLayer with the name
+NEOML_API CProblemSourceLayer* ProblemSource( CDnn& dnn, const char* name,
+	TBlobType labelType, int batchSize, const CPtr<const IProblem>& problem, bool shuffle = false, unsigned seed = 42 );
+
+//---------------------------------------------------------------------------------------------------------------------
 
 // CDnnModelWrapper is the base class wrapping the trained neural network into the IModel interface
 class NEOML_API CDnnModelWrapper : public IModel {
 public:
 	explicit CDnnModelWrapper(IMathEngine& mathEngine, unsigned int seed = 0xDEADFACE);
 
-	int GetClassCount() const override;
+	int GetClassCount() const override { return ClassCount; }
 	bool Classify(const CFloatVectorDesc& data, CClassificationResult& result) const override;
 	void Serialize(CArchive& archive) override;
 
 protected:
 	int ClassCount;
 	float SourceEmptyFill;
 	mutable CRandom Random;
-	mutable CDnn Dnn;	// the network
-	CPtr<CSourceLayer> SourceLayer;	// the reference to the source layer
-	CPtr<CSinkLayer> SinkLayer;		// the reference to the terminator layer
-	CPtr<CDnnBlob> SourceBlob;			// the source data blob
-	mutable CArray<float> tempExp;		// the temporary array for exponent values to calculate softmax
+	mutable CDnn Dnn; // the network
+	CPtr<CSourceLayer> SourceLayer; // the reference to the source layer
+	CPtr<CSinkLayer> SinkLayer; // the reference to the terminator layer
+	CPtr<CDnnBlob> SourceBlob; // the source data blob
+	mutable CArray<float> tempExp; // the temporary array for exponent values to calculate softmax
 
 	static const char* const SourceLayerName;
 	static const char* const SinkLayerName;
@@ -101,7 +125,7 @@ class NEOML_API CDnnModelWrapper : public IModel {
 	bool classify( CClassificationResult& result ) const;
 };
 
-///////////////////////////////////////////////////////////////////////////////////////////////////////
+//---------------------------------------------------------------------------------------------------------------------
 
 // CDnnTrainingModelWrapper is the base class wrapping the neural network 
 // into an ITrainingModel interface so the network can be trained using the Train method

NeoML/include/NeoML/TraditionalML/Shuffler.h

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -23,7 +23,7 @@ namespace NeoML {
 // The shuffler class
 // Uses the standard shuffling algorithm, not all at once but sequentially; as a result, the first N positions are only shuffled among themselves
 // For example, you can use it to get the random indices in an array
-class NEOML_API CShuffler {
+class NEOML_API CShuffler final {
 public:
 	CShuffler( CRandom& _random, int count );
 
@@ -34,6 +34,10 @@ class NEOML_API CShuffler {
 	int SetNext( int index );
 	// Finishes shuffling and returns all indices
 	const CArray<int>& GetAllIndices();
+	// Is shuffling finished
+	bool IsFinished() const { return nextIndex == indices.Size(); }
+	// Reset state to use shuffler again for the same array
+	void Reset() { nextIndex = 0; }
 
 private:
 	CRandom& random;

NeoML/src/Dnn/Layers/FullyConnectedSourceLayer.cpp

@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -310,15 +310,20 @@ bool CFullyConnectedSourceLayer::isBatchLoaded( int index ) const
 	return ( batchFirstLoadedIndex <= index && index <= batchLastLoadedIndex );
 }
 
-CLayerWrapper<CFullyConnectedSourceLayer> FullyConnectedSource( TBlobType labelType,
-	int batchSize, int maxBatchCount, IProblem* problem )
+// Creates CFullyConnectedSourceLayer with the name
+CFullyConnectedSourceLayer* FullyConnectedSource( CDnn& dnn, const char* name,
+	TBlobType labelType, int batchSize, int maxBatchCount, IProblem* problem, int numberOfElements, bool isZeroFreeTerm )
 {
-	return CLayerWrapper<CFullyConnectedSourceLayer>( "FullyConnectedSource", [=, &problem]( CFullyConnectedSourceLayer* result ) {
-		result->SetLabelType( labelType );
-		result->SetBatchSize( batchSize );
-		result->SetMaxBatchCount( maxBatchCount );
-		result->SetProblem( problem );
-	} );
+	CPtr<CFullyConnectedSourceLayer> result = new CFullyConnectedSourceLayer( dnn.GetMathEngine() );
+	result->SetLabelType( labelType );
+	result->SetBatchSize( batchSize );
+	result->SetMaxBatchCount( maxBatchCount );
+	result->SetProblem( problem );
+	result->SetNumberOfElements( numberOfElements );
+	result->SetZeroFreeTerm( isZeroFreeTerm );
+	result->SetName( name );
+	dnn.AddLayer( *result );
+	return result;
 }
 
 } // namespace NeoML