Add Last batch policy support for remaining image readers

rocAL/include/readers/file_source_reader.h

@@ -85,15 +85,13 @@ class FileSourceReader : public Reader {
     unsigned _curr_file_idx;
     FILE *_current_fPtr;
     unsigned _current_file_size;
-    unsigned _shard_start_idx;
     std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
     std::string _last_id;
     std::string _last_file_name, _last_file_path, _absolute_file_path;
     size_t _shard_id = 0;
     size_t _shard_count = 1;  // equivalent of batch size
     int32_t _shard_size = -1;
     size_t _batch_size = 1;
-    size_t _padded_samples = 0;
     bool _loop;
     bool _shuffle;
     int _read_counter = 0;

rocAL/include/readers/image/caffe2_lmdb_record_reader.h

@@ -68,6 +68,8 @@ class Caffe2LMDBRecordReader : public Reader {
 
     Caffe2LMDBRecordReader();
 
+    size_t last_batch_padded_size() override; // The size of the number of samples padded in the last batch
+
    private:
     //! opens the folder containnig the images
     Reader::Status Caffe2_LMDB_reader();
@@ -78,7 +80,7 @@ class Caffe2LMDBRecordReader : public Reader {
     DIR* _sub_dir;
     struct dirent* _entity;
     std::vector<std::string> _file_names;
-    std::map<std::string, unsigned int> _file_size;
+    std::map<std::string, unsigned int> _file_size, _all_shard_file_sizes_padded;
     unsigned _curr_file_idx;
     unsigned _current_file_size;
     std::string _last_id;
@@ -87,24 +89,16 @@ class Caffe2LMDBRecordReader : public Reader {
     size_t _shard_id = 0;
     size_t _shard_count = 1;  // equivalent of batch size
     bool _last_rec;
-    //!< _batch_count Defines the quantum count of the images to be read. It's usually equal to the user's batch size.
-    /// The loader will repeat images if necessary to be able to have images available in multiples of the load_batch_count,
-    /// for instance if there are 10 images in the dataset and _batch_count is 3, the loader repeats 2 images as if there are 12 images available.
-    size_t _batch_count = 1;
-    size_t _file_id = 0;
+    size_t _batch_size = 1;
     size_t _in_batch_read_count = 0;
     bool _loop;
     bool _shuffle;
     int _read_counter = 0;
     uint _file_byte_size;
     void incremenet_read_ptr();
     int release();
-    size_t get_file_shard_id();
     //!< _file_count_all_shards total_number of files in to figure out the max_batch_size (usually needed for distributed training).
     size_t _file_count_all_shards;
-    void incremenet_file_id() { _file_id++; }
-    void replicate_last_image_to_fill_last_shard();
-    void replicate_last_batch_to_pad_partial_shard();
     void read_image(unsigned char* buff, std::string file_name);
     void read_image_names();
     std::map<std::string, uint> _image_record_starting;
@@ -116,4 +110,16 @@ class Caffe2LMDBRecordReader : public Reader {
     MDB_txn* _read_mdb_txn;
     MDB_cursor* _read_mdb_cursor;
     void open_env_for_read_image();
+    int32_t _shard_size = -1;
+    ShardingInfo _sharding_info = ShardingInfo();  // The members of ShardingInfo determines how the data is distributed among the shards and how the last batch is processed by the pipeline.
+    std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
+    size_t _last_batch_padded_size = 0;
+    bool _stick_to_shard = false;
+    bool _pad_last_batch_repeated = false;
+    size_t actual_shard_size_without_padding(); // Number of files belonging to a shard (without padding)
+    size_t largest_shard_size_without_padding(); // Number of files belonging to a shard (with padding)
+    //!< Used to advance to the next shard's data to increase the entropy of the data seen by the pipeline>
+    void increment_shard_id();
+    void increment_curr_file_idx();
+    void compute_start_and_end_idx_of_all_shards();     // Start Idx of all the Shards
 };

rocAL/include/readers/image/caffe_lmdb_record_reader.h

@@ -68,6 +68,8 @@ class CaffeLMDBRecordReader : public Reader {
 
     CaffeLMDBRecordReader();
 
+    size_t last_batch_padded_size() override; // The size of the number of samples padded in the last batch
+
    private:
     //! opens the folder containnig the images
     Reader::Status folder_reading();
@@ -85,10 +87,7 @@ class CaffeLMDBRecordReader : public Reader {
     size_t _shard_id = 0;
     size_t _shard_count = 1;  // equivalent of batch size
     bool _last_rec;
-    //!< _batch_count Defines the quantum count of the images to be read. It's usually equal to the user's batch size.
-    /// The loader will repeat images if necessary to be able to have images available in multiples of the load_batch_count,
-    /// for instance if there are 10 images in the dataset and _batch_count is 3, the loader repeats 2 images as if there are 12 images available.
-    size_t _batch_count = 1;
+    size_t _batch_size = 1;
     size_t _file_id = 0;
     size_t _in_batch_read_count = 0;
     bool _loop;
@@ -102,17 +101,25 @@ class CaffeLMDBRecordReader : public Reader {
     uint _file_byte_size;
     void incremenet_read_ptr();
     int release();
-    size_t get_file_shard_id();
     //!< _file_count_all_shards total_number of files in to figure out the max_batch_size (usually needed for distributed training).
     size_t _file_count_all_shards;
-    void incremenet_file_id() { _file_id++; }
-    void replicate_last_image_to_fill_last_shard();
-    void replicate_last_batch_to_pad_partial_shard();
     void read_image(unsigned char* buff, std::string _file_name);
     void read_image_names();
     std::map<std::string, uint> _image_record_starting;
     int _open_env = 1;
     int rc;
     void open_env_for_read_image();
     std::shared_ptr<MetaDataReader> _meta_data_reader = nullptr;
+    int32_t _shard_size = -1;
+    std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
+    size_t _last_batch_padded_size = 0;
+    ShardingInfo _sharding_info = ShardingInfo();  // The members of ShardingInfo determines how the data is distributed among the shards and how the last batch is processed by the pipeline.
+    bool _stick_to_shard = false;
+    bool _pad_last_batch_repeated = false;
+    size_t actual_shard_size_without_padding(); // Number of files belonging to a shard (without padding)
+    size_t largest_shard_size_without_padding(); // Number of files belonging to a shard (with padding)
+    //!< Used to advance to the next shard's data to increase the entropy of the data seen by the pipeline>
+    void increment_shard_id();
+    void increment_curr_file_idx();
+    void compute_start_and_end_idx_of_all_shards();     // Start Idx of all the Shards
 };

rocAL/include/readers/image/cifar10_data_reader.h

@@ -64,6 +64,8 @@ class CIFAR10DataReader : public Reader {
 
     unsigned get_file_index() { return _last_file_idx; }
 
+    size_t last_batch_padded_size() override; // The size of the number of samples padded in the last batch
+
    private:
     //! opens the folder containing the images
     Reader::Status open_folder();
@@ -87,15 +89,26 @@ class CIFAR10DataReader : public Reader {
     //!< _raw_file_size of each file to read
     const size_t _raw_file_size = (32 * 32 * 3 + 1);  // todo:: need to add an option in reader config to take this.
     size_t _total_file_size;
-    //!< _batch_count Defines the quantum count of the images to be read. It's usually equal to the user's batch size.
-    /// The loader will repeat images if necessary to be able to have images available in multiples of the load_batch_count,
-    /// for instance if there are 10 images in the dataset and _batch_count is 3, the loader repeats 2 images as if there are 12 images available.
-    size_t _batch_count = 1;
-    size_t _file_id = 0;
-    size_t _in_batch_read_count = 0;
+    size_t _batch_size = 1;
     bool _loop;
+    bool _shuffle;
     int _read_counter = 0;
     void incremenet_read_ptr();
     int release();
-    void incremenet_file_id() { _file_id++; }
+    int32_t _shard_size = -1;
+    size_t _shard_id = 0;
+    size_t _shard_count = 1;
+    std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
+    //!< _file_count_all_shards total_number of files in to figure out the max_batch_size (usually needed for distributed training).
+    size_t _file_count_all_shards;
+    ShardingInfo _sharding_info = ShardingInfo();  // The members of ShardingInfo determines how the data is distributed among the shards and how the last batch is processed by the pipeline.
+    size_t _last_batch_padded_size = 0;
+    bool _stick_to_shard = false;
+    bool _pad_last_batch_repeated = false;
+    size_t actual_shard_size_without_padding(); // Number of files belonging to a shard (without padding)
+    size_t largest_shard_size_without_padding(); // Number of files belonging to a shard (with padding)
+    //!< Used to advance to the next shard's data to increase the entropy of the data seen by the pipeline>
+    void increment_shard_id();
+    void increment_curr_file_idx();
+    void compute_start_and_end_idx_of_all_shards();
 };

rocAL/include/readers/image/coco_file_source_reader.h

@@ -66,6 +66,8 @@ class COCOFileSourceReader : public Reader {
 
     COCOFileSourceReader();
 
+    size_t last_batch_padded_size() override; // The size of the number of samples padded in the last batch
+
    private:
     std::shared_ptr<MetaDataReader> _meta_data_reader = nullptr;
     //! opens the folder containnig the images
@@ -85,23 +87,26 @@ class COCOFileSourceReader : public Reader {
     std::string _last_id;
     std::string _last_file_name;
     size_t _shard_id = 0;
-    size_t _shard_count = 1;  // equivalent of batch size
-    //!< _batch_count Defines the quantum count of the images to be read. It's usually equal to the user's batch size.
-    /// The loader will repeat images if necessary to be able to have images available in multiples of the load_batch_count,
-    /// for instance if there are 10 images in the dataset and _batch_count is 3, the loader repeats 2 images as if there are 12 images available.
-    size_t _batch_count = 1;
-    size_t _file_id = 0;
-    size_t _in_batch_read_count = 0;
+    size_t _shard_count = 1;
+    size_t _batch_size = 1;
     bool _loop;
     bool _shuffle;
     int _read_counter = 0;
     //!< _file_count_all_shards total_number of files in to figure out the max_batch_size (usually needed for distributed training).
     size_t _file_count_all_shards;
     void incremenet_read_ptr();
     int release();
-    size_t get_file_shard_id();
-    void incremenet_file_id() { _file_id++; }
-    void replicate_last_image_to_fill_last_shard();
-    void replicate_last_batch_to_pad_partial_shard();
     void shuffle_with_aspect_ratios();
+    void increment_curr_file_idx();
+    ShardingInfo _sharding_info = ShardingInfo();  // The members of ShardingInfo determines how the data is distributed among the shards and how the last batch is processed by the pipeline.
+    size_t _last_batch_padded_size = 0;
+    bool _stick_to_shard = false;
+    bool _pad_last_batch_repeated = false;
+    int32_t _shard_size = -1;
+    std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
+    size_t actual_shard_size_without_padding(); // Number of files belonging to a shard (without padding)
+    size_t largest_shard_size_without_padding(); // Number of files belonging to a shard (with padding)
+    //!< Used to advance to the next shard's data to increase the entropy of the data seen by the pipeline>
+    void increment_shard_id();
+    void compute_start_and_end_idx_of_all_shards();     // Start Idx of all the Shards
 };

rocAL/include/readers/image/mxnet_recordio_reader.h

@@ -66,8 +66,10 @@ class MXNetRecordIOReader : public Reader {
 
     MXNetRecordIOReader();
 
+    size_t last_batch_padded_size() override; // The size of the number of samples padded in the last batch
+
    private:
-    //! opens the folder containnig the images
+    //! opens the folder containig the images
     Reader::Status record_reading();
     Reader::Status MXNet_reader();
     std::string _path;
@@ -83,11 +85,8 @@ class MXNetRecordIOReader : public Reader {
     int64_t _last_seek_pos;
     int64_t _last_data_size;
     size_t _shard_id = 0;
-    size_t _shard_count = 1;  // equivalent of batch size
-    //!< _batch_count Defines the quantum count of the images to be read. It's usually equal to the user's batch size.
-    /// The loader will repeat images if necessary to be able to have images available in multiples of the load_batch_count,
-    /// for instance if there are 10 images in the dataset and _batch_count is 3, the loader repeats 2 images as if there are 12 images available.
-    size_t _batch_count = 1;
+    size_t _shard_count = 1;
+    size_t _batch_size = 1;
     size_t _file_id = 0;
     size_t _in_batch_read_count = 0;
     bool _loop;
@@ -97,10 +96,6 @@ class MXNetRecordIOReader : public Reader {
     size_t _file_count_all_shards;
     void incremenet_read_ptr();
     int release();
-    size_t get_file_shard_id();
-    void incremenet_file_id() { _file_id++; }
-    void replicate_last_image_to_fill_last_shard();
-    void replicate_last_batch_to_pad_partial_shard();
     void read_image(unsigned char* buff, int64_t seek_position, int64_t data_size);
     void read_image_names();
     uint32_t DecodeFlag(uint32_t rec) { return (rec >> 29U) & 7U; };
@@ -110,4 +105,16 @@ class MXNetRecordIOReader : public Reader {
     const uint32_t _kMagic = 0xced7230a;
     int64_t _seek_pos, _data_size_to_read;
     ImageRecordIOHeader _hdr;
+    int32_t _shard_size = -1;
+    ShardingInfo _sharding_info = ShardingInfo();  // The members of ShardingInfo determines how the data is distributed among the shards and how the last batch is processed by the pipeline.
+    std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
+    bool _stick_to_shard = false;
+    bool _pad_last_batch_repeated = false;
+    size_t _last_batch_padded_size = 0;
+    size_t actual_shard_size_without_padding(); // Number of files belonging to a shard (without padding)
+    size_t largest_shard_size_without_padding(); // Number of files belonging to a shard (with padding)
+    //!< Used to advance to the next shard's data to increase the entropy of the data seen by the pipeline>
+    void increment_shard_id();
+    void increment_curr_file_idx();
+    void compute_start_and_end_idx_of_all_shards();     // Start Idx of all the Shards
 };

rocAL/include/readers/image/tf_record_reader.h

@@ -68,6 +68,8 @@ class TFRecordReader : public Reader {
 
     TFRecordReader();
 
+    size_t last_batch_padded_size() override; // The size of the number of samples padded in the last batch
+
    private:
     //! opens the folder containnig the images
     Reader::Status tf_record_reader();
@@ -81,7 +83,7 @@ class TFRecordReader : public Reader {
     DIR *_sub_dir;
     struct dirent *_entity;
     std::vector<std::string> _file_names;
-    std::map<std::string, unsigned int> _file_size;
+    std::map<std::string, unsigned int> _file_size, _all_shard_file_sizes_padded;
     unsigned _curr_file_idx;
     unsigned _current_file_size;
     std::string _last_id;
@@ -90,12 +92,8 @@ class TFRecordReader : public Reader {
     size_t _shard_id = 0;
     size_t _shard_count = 1;  // equivalent of batch size
     bool _last_rec;
-    //!< _batch_count Defines the quantum count of the images to be read. It's usually equal to the user's batch size.
-    /// The loader will repeat images if necessary to be able to have images available in multiples of the load_batch_count,
-    /// for instance if there are 10 images in the dataset and _batch_count is 3, the loader repeats 2 images as if there are 12 images available.
-    size_t _batch_count = 1;
+    size_t _batch_size = 1;
     size_t _file_id = 0;
-    size_t _in_batch_read_count = 0;
     bool _loop;
     bool _shuffle;
     int _read_counter = 0;
@@ -108,11 +106,19 @@ class TFRecordReader : public Reader {
     tensorflow::Feature _single_feature;
     void incremenet_read_ptr();
     int release();
-    size_t get_file_shard_id();
-    void incremenet_file_id() { _file_id++; }
-    void replicate_last_image_to_fill_last_shard();
-    void replicate_last_batch_to_pad_partial_shard();
     Reader::Status read_image(unsigned char *buff, std::string record_file_name, uint file_size);
     Reader::Status read_image_names(std::ifstream &file_contents, uint file_size);
     std::map<std::string, uint> _image_record_starting;
+    ShardingInfo _sharding_info = ShardingInfo();  // The members of ShardingInfo determines how the data is distributed among the shards and how the last batch is processed by the pipeline.
+    size_t _last_batch_padded_size = 0;
+    bool _stick_to_shard = false;
+    bool _pad_last_batch_repeated = false;
+    int32_t _shard_size = -1;
+    std::vector<unsigned> _shard_start_idx_vector, _shard_end_idx_vector;
+    void increment_curr_file_idx();
+    size_t actual_shard_size_without_padding(); // Number of files belonging to a shard (without padding)
+    size_t largest_shard_size_without_padding(); // Number of files belonging to a shard (with padding)
+    //!< Used to advance to the next shard's data to increase the entropy of the data seen by the pipeline>
+    void compute_start_and_end_idx_of_all_shards();     // Start Idx of all the Shards
+    void increment_shard_id();
 };

rocAL/source/readers/file_source_reader.cpp

@@ -260,9 +260,10 @@ Reader::Status FileSourceReader::generate_file_names() {
         ERR("FileReader ShardID [" + TOSTR(_shard_id) + "] Did not load any file from " + _folder_path)
 
     auto dataset_size = _file_count_all_shards;
+    size_t padded_samples = 0;
     // Pad the _file_names with last element of the shard in the vector when _pad_last_batch_repeated is True
-    _padded_samples = ((_shard_size > 0) ? _shard_size : largest_shard_size_without_padding()) % _batch_size;
-    _last_batch_padded_size = ((_batch_size > 1) && (_padded_samples > 0 )) ? (_batch_size - _padded_samples) : 0;
+    padded_samples = ((_shard_size > 0) ? _shard_size : largest_shard_size_without_padding()) % _batch_size;
+    _last_batch_padded_size = ((_batch_size > 1) && (padded_samples > 0)) ? (_batch_size - padded_samples) : 0;
 
     if (_pad_last_batch_repeated == true) {
                                             // pad the last sample when the dataset_size is not divisible by