Merge pull request official-stockfish#259 from Sopel97/builtin_interleave

Add support for multiple datasets in the trainer. The datasets are in…

Author: vondele

lib/nnue_training_data_formats.h

@@ -6726,6 +6726,12 @@ namespace binpack
         }
     }
 
+    inline std::ifstream::pos_type filesize(const char* filename)
+    {
+        std::ifstream in(filename, std::ifstream::ate | std::ifstream::binary);
+        return in.tellg();
+    }
+
     struct CompressedTrainingDataFile
     {
         struct Header
@@ -6737,12 +6743,15 @@ namespace binpack
             m_path(std::move(path)),
             m_file(m_path, std::ios_base::binary | std::ios_base::in | std::ios_base::out | om)
         {
+            // Racey but who cares
+            m_sizeBytes = filesize(m_path.c_str());
         }
 
         void append(const char* data, std::uint32_t size)
         {
             writeChunkHeader({size});
             m_file.write(data, size);
+            m_sizeBytes += size + 8;
         }
 
         [[nodiscard]] bool hasNextChunk()
@@ -6756,6 +6765,11 @@ namespace binpack
             return !m_file.eof();
         }
 
+        void seek_to_start()
+        {
+            m_file.seekg(0);
+        }
+
         [[nodiscard]] std::vector<unsigned char> readNextChunk()
         {
             auto size = readChunkHeader().chunkSize;
@@ -6764,9 +6778,15 @@ namespace binpack
             return data;
         }
 
+        [[nodiscard]] std::size_t sizeBytes() const
+        {
+            return m_sizeBytes;
+        }
+
     private:
         std::string m_path;
         std::fstream m_file;
+        std::size_t m_sizeBytes;
 
         void writeChunkHeader(Header h)
         {
@@ -7558,21 +7578,32 @@ namespace binpack
 
         CompressedTrainingDataEntryParallelReader(
             int concurrency,
-            std::string path,
+            std::vector<std::string> paths,
             std::ios_base::openmode om = std::ios_base::app,
+            bool cyclic = false,
             std::function<bool(const TrainingDataEntry&)> skipPredicate = nullptr
         ) :
             m_concurrency(concurrency),
-            m_inputFile(path, om),
             m_bufferOffset(0),
+            m_cyclic(cyclic),
             m_skipPredicate(std::move(skipPredicate))
         {
             m_numRunningWorkers.store(0);
-            if (!m_inputFile.hasNextChunk())
+            std::vector<double> sizes; // discrete distribution wants double weights
+            for (const auto& path : paths)
             {
-                return;
+                auto& file = m_inputFiles.emplace_back(path, om);
+
+                if (!file.hasNextChunk())
+                {
+                    return;
+                }
+
+                sizes.emplace_back(static_cast<double>(file.sizeBytes()));
             }
 
+            m_inputFileDistribution = std::discrete_distribution<>(sizes.begin(), sizes.end());
+
             m_stopFlag.store(false);
 
             auto worker = [this]()
@@ -7742,8 +7773,10 @@ namespace binpack
 
     private:
         int m_concurrency;
-        CompressedTrainingDataFile m_inputFile;
+        std::vector<CompressedTrainingDataFile> m_inputFiles;
+        std::discrete_distribution<> m_inputFileDistribution;
         std::atomic_int m_numRunningWorkers;
+        bool m_cyclic;
 
         static constexpr int threadBufferSize = 256 * 256 * 16;
 
@@ -7763,17 +7796,24 @@ namespace binpack
         {
             if (m_offset + sizeof(PackedTrainingDataEntry) + 2 > m_chunk.size())
             {
+                auto& prng = rng::get_thread_local_rng();
+                const std::size_t fileId = m_inputFileDistribution(prng);
+                auto& inputFile = m_inputFiles[fileId];
+
                 std::unique_lock lock(m_fileMutex);
 
-                if (!m_inputFile.hasNextChunk())
-                {
-                    return true;
-                }
-                else
+                if (!inputFile.hasNextChunk())
                 {
-                    m_chunk = m_inputFile.readNextChunk();
-                    m_offset = 0;
+                    if (m_cyclic)
+                    {
+                        inputFile.seek_to_start();
+                    }
+                    else
+                        return true;
                 }
+
+                m_chunk = inputFile.readNextChunk();
+                m_offset = 0;
             }
 
             return false;

lib/nnue_training_data_stream.h

@@ -183,9 +183,9 @@ namespace training_data {
         static constexpr auto openmode = std::ios::in | std::ios::binary;
         static inline const std::string extension = "binpack";
 
-        BinpackSfenInputParallelStream(int concurrency, std::string filename, bool cyclic, std::function<bool(const TrainingDataEntry&)> skipPredicate) :
-            m_stream(std::make_unique<binpack::CompressedTrainingDataEntryParallelReader>(concurrency, filename, openmode, skipPredicate)),
-            m_filename(filename),
+        BinpackSfenInputParallelStream(int concurrency, const std::vector<std::string>& filenames, bool cyclic, std::function<bool(const TrainingDataEntry&)> skipPredicate) :
+            m_stream(std::make_unique<binpack::CompressedTrainingDataEntryParallelReader>(concurrency, filenames, openmode, cyclic, skipPredicate)),
+            m_filenames(filenames),
             m_concurrency(concurrency),
             m_eof(false),
             m_cyclic(cyclic),
@@ -199,12 +199,6 @@ namespace training_data {
             auto v = m_stream->next();
             if (!v.has_value())
             {
-                if (m_cyclic)
-                {
-                    m_stream = std::make_unique<binpack::CompressedTrainingDataEntryParallelReader>(m_concurrency, m_filename, openmode, m_skipPredicate);
-                    return m_stream->next();
-                }
-
                 m_eof = true;
                 return std::nullopt;
             }
@@ -217,32 +211,7 @@ namespace training_data {
             auto k = m_stream->fill(v, n);
             if (n != k)
             {
-                if (m_cyclic)
-                {
-                    m_stream = std::make_unique<binpack::CompressedTrainingDataEntryParallelReader>(m_concurrency, m_filename, openmode, m_skipPredicate);
-                    n -= k;
-                    k = m_stream->fill(v, n);
-                    if (k == 0)
-                    {
-                        // No data in the file
-                        m_eof = true;
-                        return;
-                    }
-                    else if (k == n)
-                    {
-                        // We're done
-                        return;
-                    }
-                    else
-                    {
-                        // We need to read again
-                        this->fill(v, n - k);
-                    }
-                }
-                else
-                {
-                    m_eof = true;
-                }
+                m_eof = true;
             }
         }
 
@@ -255,7 +224,7 @@ namespace training_data {
 
     private:
         std::unique_ptr<binpack::CompressedTrainingDataEntryParallelReader> m_stream;
-        std::string m_filename;
+        std::vector<std::string> m_filenames;
         int m_concurrency;
         bool m_eof;
         bool m_cyclic;
@@ -272,13 +241,13 @@ namespace training_data {
         return nullptr;
     }
 
-    inline std::unique_ptr<BasicSfenInputStream> open_sfen_input_file_parallel(int concurrency, const std::string& filename, bool cyclic, std::function<bool(const TrainingDataEntry&)> skipPredicate = nullptr)
+    inline std::unique_ptr<BasicSfenInputStream> open_sfen_input_file_parallel(int concurrency, const std::vector<std::string>& filenames, bool cyclic, std::function<bool(const TrainingDataEntry&)> skipPredicate = nullptr)
     {
         // TODO (low priority): optimize and parallelize .bin reading.
-        if (has_extension(filename, BinSfenInputStream::extension))
-            return std::make_unique<BinSfenInputStream>(filename, cyclic, std::move(skipPredicate));
-        else if (has_extension(filename, BinpackSfenInputParallelStream::extension))
-            return std::make_unique<BinpackSfenInputParallelStream>(concurrency, filename, cyclic, std::move(skipPredicate));
+        if (has_extension(filenames[0], BinSfenInputStream::extension))
+            return std::make_unique<BinSfenInputStream>(filenames[0], cyclic, std::move(skipPredicate));
+        else if (has_extension(filenames[0], BinpackSfenInputParallelStream::extension))
+            return std::make_unique<BinpackSfenInputParallelStream>(concurrency, filenames, cyclic, std::move(skipPredicate));
 
         return nullptr;
     }

nnue_dataset.py

@@ -67,13 +67,18 @@ def get_fens(self):
         return strings
 
 FenBatchPtr = ctypes.POINTER(FenBatch)
-# EXPORT FenBatchStream* CDECL create_fen_batch_stream(int concurrency, const char* filename, int batch_size, bool cyclic, bool filtered, int random_fen_skipping, bool wld_filtered, int param_index)
+# EXPORT FenBatchStream* CDECL create_fen_batch_stream(int concurrency, int num_files, const char* const* filenames, int batch_size, bool cyclic, bool filtered, int random_fen_skipping, bool wld_filtered, int early_fen_skipping, int param_index)
 create_fen_batch_stream = dll.create_fen_batch_stream
 create_fen_batch_stream.restype = ctypes.c_void_p
-create_fen_batch_stream.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_bool, ctypes.c_bool, ctypes.c_int, ctypes.c_bool, ctypes.c_int, ctypes.c_int]
+create_fen_batch_stream.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_char_p), ctypes.c_int, ctypes.c_bool, ctypes.c_bool, ctypes.c_int, ctypes.c_bool, ctypes.c_int, ctypes.c_int]
 destroy_fen_batch_stream = dll.destroy_fen_batch_stream
 destroy_fen_batch_stream.argtypes = [ctypes.c_void_p]
 
+def make_fen_batch_stream(concurrency, filenames, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index):
+    filenames_ = (ctypes.c_char_p * len(filenames))()
+    filenames_[:] = [filename.encode('utf-8') for filename in filenames]
+    return create_fen_batch_stream(concurrency, len(filenames), filenames_, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
+
 fetch_next_fen_batch = dll.fetch_next_fen_batch
 fetch_next_fen_batch.restype = FenBatchPtr
 fetch_next_fen_batch.argtypes = [ctypes.c_void_p]
@@ -103,9 +108,9 @@ def __init__(
         self.param_index = param_index
 
         if batch_size:
-            self.stream = create_fen_batch_stream(self.num_workers, self.filename, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
+            self.stream = make_fen_batch_stream(self.num_workers, [self.filename], batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
         else:
-            self.stream = create_fen_batch_stream(self.num_workers, self.filename, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
+            self.stream = make_fen_batch_stream(self.num_workers, [self.filename], cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
 
     def __iter__(self):
         return self
@@ -131,7 +136,7 @@ def __init__(
         destroy_stream,
         fetch_next,
         destroy_part,
-        filename,
+        filenames,
         cyclic,
         num_workers,
         batch_size=None,
@@ -147,7 +152,7 @@ def __init__(
         self.destroy_stream = destroy_stream
         self.fetch_next = fetch_next
         self.destroy_part = destroy_part
-        self.filename = filename.encode('utf-8')
+        self.filenames = filenames
         self.cyclic = cyclic
         self.num_workers = num_workers
         self.batch_size = batch_size
@@ -158,9 +163,9 @@ def __init__(
         self.device = device
 
         if batch_size:
-            self.stream = self.create_stream(self.feature_set, self.num_workers, self.filename, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
+            self.stream = self.create_stream(self.feature_set, self.num_workers, self.filenames, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
         else:
-            self.stream = self.create_stream(self.feature_set, self.num_workers, self.filename, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
+            self.stream = self.create_stream(self.feature_set, self.num_workers, self.filenames, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
 
     def __iter__(self):
         return self
@@ -178,14 +183,19 @@ def __next__(self):
     def __del__(self):
         self.destroy_stream(self.stream)
 
-#    EXPORT Stream<SparseBatch>* CDECL create_sparse_batch_stream(const char* feature_set_c, int concurrency, const char* filename, int batch_size, bool cyclic,
+#    EXPORT Stream<SparseBatch>* CDECL create_sparse_batch_stream(const char* feature_set_c, int concurrency, int num_files, const char* const* filenames, int batch_size, bool cyclic,
 #                                                                 bool filtered, int random_fen_skipping, bool wld_filtered, int early_fen_skipping, int param_index)
 create_sparse_batch_stream = dll.create_sparse_batch_stream
 create_sparse_batch_stream.restype = ctypes.c_void_p
-create_sparse_batch_stream.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_bool, ctypes.c_bool, ctypes.c_int, ctypes.c_bool, ctypes.c_int, ctypes.c_int]
+create_sparse_batch_stream.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_char_p), ctypes.c_int, ctypes.c_bool, ctypes.c_bool, ctypes.c_int, ctypes.c_bool, ctypes.c_int, ctypes.c_int]
 destroy_sparse_batch_stream = dll.destroy_sparse_batch_stream
 destroy_sparse_batch_stream.argtypes = [ctypes.c_void_p]
 
+def make_sparse_batch_stream(feature_set, concurrency, filenames, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index):
+    filenames_ = (ctypes.c_char_p * len(filenames))()
+    filenames_[:] = [filename.encode('utf-8') for filename in filenames]
+    return create_sparse_batch_stream(feature_set, concurrency, len(filenames), filenames_, batch_size, cyclic, filtered, random_fen_skipping, wld_filtered, early_fen_skipping, param_index)
+
 fetch_next_sparse_batch = dll.fetch_next_sparse_batch
 fetch_next_sparse_batch.restype = SparseBatchPtr
 fetch_next_sparse_batch.argtypes = [ctypes.c_void_p]
@@ -211,14 +221,14 @@ def make_sparse_batch_from_fens(feature_set, fens, scores, plies, results):
     return b
 
 class SparseBatchProvider(TrainingDataProvider):
-    def __init__(self, feature_set, filename, batch_size, cyclic=True, num_workers=1, filtered=False, random_fen_skipping=0, wld_filtered=False, early_fen_skipping=-1, param_index=0, device='cpu'):
+    def __init__(self, feature_set, filenames, batch_size, cyclic=True, num_workers=1, filtered=False, random_fen_skipping=0, wld_filtered=False, early_fen_skipping=-1, param_index=0, device='cpu'):
         super(SparseBatchProvider, self).__init__(
             feature_set,
-            create_sparse_batch_stream,
+            make_sparse_batch_stream,
             destroy_sparse_batch_stream,
             fetch_next_sparse_batch,
             destroy_sparse_batch,
-            filename,
+            filenames,
             cyclic,
             num_workers,
             batch_size,
@@ -230,10 +240,10 @@ def __init__(self, feature_set, filename, batch_size, cyclic=True, num_workers=1
             device)
 
 class SparseBatchDataset(torch.utils.data.IterableDataset):
-  def __init__(self, feature_set, filename, batch_size, cyclic=True, num_workers=1, filtered=False, random_fen_skipping=0, wld_filtered=False, early_fen_skipping=-1, param_index=0, device='cpu'):
+  def __init__(self, feature_set, filenames, batch_size, cyclic=True, num_workers=1, filtered=False, random_fen_skipping=0, wld_filtered=False, early_fen_skipping=-1, param_index=0, device='cpu'):
     super(SparseBatchDataset).__init__()
     self.feature_set = feature_set
-    self.filename = filename
+    self.filenames = filenames
     self.batch_size = batch_size
     self.cyclic = cyclic
     self.num_workers = num_workers
@@ -245,7 +255,7 @@ def __init__(self, feature_set, filename, batch_size, cyclic=True, num_workers=1
     self.device = device
 
   def __iter__(self):
-    return SparseBatchProvider(self.feature_set, self.filename, self.batch_size, cyclic=self.cyclic, num_workers=self.num_workers,
+    return SparseBatchProvider(self.feature_set, self.filenames, self.batch_size, cyclic=self.cyclic, num_workers=self.num_workers,
                                filtered=self.filtered, random_fen_skipping=self.random_fen_skipping, wld_filtered=self.wld_filtered, early_fen_skipping = self.early_fen_skipping, param_index=self.param_index, device=self.device)
 
 class FixedNumBatchesDataset(Dataset):