[NeoML] CProblemSourceLayer uses new CShuffler

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

Author: favorart

NeoML/src/Dnn/Layers/ModelWrapperLayer.cpp

@@ -26,9 +26,8 @@ limitations under the License.
 namespace NeoML {
 
 // Shuffles the elements array.
-static void shuffle( CArray<int>& elements, unsigned seed )
+static void shuffle( CArray<int>& elements, CRandom& random )
 {
-	CRandom random( seed );
 	CShuffler indexGenerator( random, elements.Size() );
 	CArray<int> oldElements;
 	elements.CopyTo( oldElements );
@@ -42,41 +41,26 @@ static void shuffle( CArray<int>& elements, unsigned seed )
 // Shuffles the elements of an array and returns them one by one.
 // Unlike CShuffler, it returns the elements of the array, not the indices of the elements, 
 // and also supports cyclicity (when the end of the sequence is reached, it will be shuffled again).
-class CShuffledGenerator final {
+template <typename T>
+class CShuffledElements final {
 public:
-	CShuffledGenerator( const CArray<int>& _elements, unsigned seed );
-	
-	CRandom& Random() { return random; }
+	CShuffledElements( CArray<int>&& _elements, CRandom& _random ) :
+		random( _random ), elements( std::move( _elements ) ), shuffler( random, elements.Size() ) {}
+	CShuffledElements( CShuffledElements&& ) = default;
+
+	CShuffledElements& operator=( CShuffledElements&& ) = default;
+
 	// The number of elements in the elements array.
 	int Size() const { return elements.Size(); }
 	// Generates the next element of the sequence.
-	int Next();
+	T Next() { if( shuffler.IsFinished() ) { shuffler.Reset(); } return elements[shuffler.Next()]; }
 
 private:
-	CRandom random;
-	CArray<int> elements; // The shuffled elements array.
-	int position = 0; // The position in the elements array.
+	CRandom& random; // The random generator
+	CArray<T> elements; // The shuffled elements array
+	CShuffler shuffler; // The index shuffled generator
 };
 
-CShuffledGenerator::CShuffledGenerator( const CArray<int>& _elements, unsigned seed ) :
-	random( seed )
-{
-	_elements.CopyTo( elements );
-	shuffle( elements, seed );
-}
-
-int CShuffledGenerator::Next()
-{
-	const int result = elements[position++];
-	NeoAssert( position <= Size() );
-	// reached the end - mix the elements.
-	if( position == Size() ) {
-		shuffle( elements, random.Next() );
-		position = 0;
-	}
-	return result;
-}
-
 //---------------------------------------------------------------------------------------------------------------------
 
 // CBalancedPairBatchGenerator
@@ -100,15 +84,28 @@ class CBalancedPairBatchGenerator : public IShuffledBatchGenerator {
 	void DeleteUnseenElement( int index ) override { unseenElementsIndices.Delete( index ); }
 
 private:
+	CRandom random;
 	// The dictionary of "label -> generator of elements indexes with this label".
-	CMap<int, CPtrOwner<CShuffledGenerator>> labelToIndexGenerator;
+	CMap<int, CShuffledElements<int>> labelToIndexGenerator;
 	// The labels generator
-	CPtrOwner<CShuffledGenerator> labelGenerator;
+	CShuffledElements<int> labelGenerator;
 	// Indexes of elements, which aren't in any batch, to detect epoch's end
 	CHashTable<int> unseenElementsIndices;
+
+	CArray<int> getLabels( const IProblem& );
 };
 
-CBalancedPairBatchGenerator::CBalancedPairBatchGenerator( const IProblem& problem, unsigned seed )
+CBalancedPairBatchGenerator::CBalancedPairBatchGenerator( const IProblem& problem, unsigned seed ) :
+	random( seed ),
+	labelGenerator( getLabels( problem ), random )
+{
+	unseenElementsIndices.SetBufferSize( problem.GetVectorCount() );
+	for( int i = 0; i < problem.GetVectorCount(); ++i ) {
+		unseenElementsIndices.Add( i );
+	}
+}
+
+CArray<int> CBalancedPairBatchGenerator::getLabels( const IProblem& problem )
 {
 	CMap<int, CArray<int>> labelToIndexes;
 	for( int i = 0; i < problem.GetVectorCount(); ++i ) {
@@ -118,16 +115,10 @@ CBalancedPairBatchGenerator::CBalancedPairBatchGenerator( const IProblem& proble
 	CArray<int> labelsUnique;
 	labelsUnique.SetBufferSize( labelToIndexes.Size() );
 	for( auto& item : labelToIndexes ) {
-		labelToIndexGenerator.CreateNewValue( item.Key ) = new CShuffledGenerator( item.Value, seed );
+		labelToIndexGenerator.Add( item.Key, CShuffledElements<int>( std::move( item.Value ), random ) );
 		labelsUnique.Add( item.Key );
 	}
-
-	labelGenerator = new CShuffledGenerator( labelsUnique, seed );
-
-	unseenElementsIndices.SetBufferSize( problem.GetVectorCount() );
-	for( int i = 0; i < problem.GetVectorCount(); ++i ) {
-		unseenElementsIndices.Add( i );
-	}
+	return labelsUnique;
 }
 
 CArray<int> CBalancedPairBatchGenerator::GenerateBatchIndexes( int batchSize )
@@ -143,12 +134,12 @@ CArray<int> CBalancedPairBatchGenerator::GenerateBatchIndexes( int batchSize )
 	const int numOfSingleClassSamples = static_cast<int>( round( idealNumOfSingleClassSamples ) );
 
 	// Sample a random class and collect numOfSingleClassSamples elements from it into a batch.
-	const int majorLabel = labelGenerator->Next();
-	NeoAssert( numOfSingleClassSamples <= labelToIndexGenerator.Get( majorLabel )->Size() );
+	const int majorLabel = labelGenerator.Next();
+	NeoAssert( numOfSingleClassSamples <= labelToIndexGenerator.Get( majorLabel ).Size() );
 	// Sample a random class and collect numOfSingleClassSamples elements from it into a batch.
 	CArray<int> batchIndexes;
 	for( int i = 0; i < numOfSingleClassSamples; ++i ) {
-		batchIndexes.Add( labelToIndexGenerator.Get( majorLabel )->Next() );
+		batchIndexes.Add( labelToIndexGenerator.Get( majorLabel ).Next() );
 	}
 
 	// The number of remaining elements, also the number of remaining classes, because for the remaining elements
@@ -164,18 +155,18 @@ CArray<int> CBalancedPairBatchGenerator::GenerateBatchIndexes( int batchSize )
 	// Will never loop if numOfOtherClasses + 1 <= totalNumOfClasses, because the generator shuffles classes without duplicates.
 	while( usedClasses.Size() < numOfOtherClasses + 1 ) {
 		// Sample the class.
-		const int label = labelGenerator->Next();
+		const int label = labelGenerator.Next();
 		// Skip the class if it has already been sampled.
 		if( usedClasses.Has( label ) ) {
 			continue;
 		}
 		usedClasses.Add( label );
 		// Sample an element from this class.
-		const int negativeSample = labelToIndexGenerator.Get( label )->Next();
+		const int negativeSample = labelToIndexGenerator.Get( label ).Next();
 		batchIndexes.Add( negativeSample );
 	}
 	// Just in case, we mix the batch before giving it to the model so that it doesn’t accidentally learn the structure of the batch.
-	shuffle( batchIndexes, labelGenerator->Random().Next() );
+	shuffle( batchIndexes, random );
 
 	NeoAssert( batchIndexes.Size() == batchSize );
 	return batchIndexes;