Added synchronized support for batches_per_epoch for quick iteration

marius-team · Jan 6, 2024 · c6c3644 · c6c3644
1 parent c04d75a
commit c6c3644
Show file tree

Hide file tree

Showing 472 changed files with 74 additions and 843 deletions.
diff --git a/.clang-format b/.clang-format
diff --git a/.flake8 b/.flake8
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
diff --git a/.github/ISSUE_TEMPLATE/documentation-improvement.md b/.github/ISSUE_TEMPLATE/documentation-improvement.md
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
diff --git a/.github/ISSUE_TEMPLATE/general-question.md b/.github/ISSUE_TEMPLATE/general-question.md
diff --git a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md
diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
diff --git a/.github/workflows/db2graph_test_postgres.yml b/.github/workflows/db2graph_test_postgres.yml
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
diff --git a/.gitignore b/.gitignore
@@ -171,3 +171,5 @@ Thumbs.db
 
 
 # End of https://www.toptal.com/developers/gitignore/api/python
+
+test_datasets
diff --git a/.gitmodules b/.gitmodules
diff --git a/CMakeLists.txt b/CMakeLists.txt
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
diff --git a/LICENSE b/LICENSE
diff --git a/MANIFEST.in b/MANIFEST.in
diff --git a/README.md b/README.md
diff --git a/docs/.nojekyll b/docs/.nojekyll
diff --git a/docs/CMakeLists.txt b/docs/CMakeLists.txt
diff --git a/docs/Doxyfile b/docs/Doxyfile
diff --git a/docs/Doxyfile.in b/docs/Doxyfile.in
diff --git a/docs/README.md b/docs/README.md
diff --git a/docs/_static/css/marius_theme.css b/docs/_static/css/marius_theme.css
diff --git a/docs/_templates/layout.html b/docs/_templates/layout.html
diff --git a/docs/assets/configuration_lp.png b/docs/assets/configuration_lp.png
diff --git a/docs/assets/configuration_nc.png b/docs/assets/configuration_nc.png
diff --git a/docs/assets/marius_arch.png b/docs/assets/marius_arch.png
diff --git a/docs/assets/samples_dismult.png b/docs/assets/samples_dismult.png
diff --git a/docs/assets/samples_gs.png b/docs/assets/samples_gs.png
diff --git a/docs/conf.py b/docs/conf.py
diff --git a/docs/config_interface/configuration.rst b/docs/config_interface/configuration.rst
diff --git a/docs/config_interface/full_schema.rst b/docs/config_interface/full_schema.rst
diff --git a/docs/config_interface/index.rst b/docs/config_interface/index.rst
diff --git a/docs/config_interface/samples.rst b/docs/config_interface/samples.rst
diff --git a/docs/db2graph/db2graph.rst b/docs/db2graph/db2graph.rst
diff --git a/docs/examples/config/index.rst b/docs/examples/config/index.rst
diff --git a/docs/examples/config/lp_custom.rst b/docs/examples/config/lp_custom.rst
diff --git a/docs/examples/config/lp_fb15k237.rst b/docs/examples/config/lp_fb15k237.rst
diff --git a/docs/examples/config/lp_paleobiology.rst b/docs/examples/config/lp_paleobiology.rst
diff --git a/docs/examples/config/nc_custom.rst b/docs/examples/config/nc_custom.rst
diff --git a/docs/examples/config/nc_ogbn_arxiv.rst b/docs/examples/config/nc_ogbn_arxiv.rst
diff --git a/docs/examples/config/resume_training.rst b/docs/examples/config/resume_training.rst
diff --git a/docs/examples/index.rst b/docs/examples/index.rst
diff --git a/docs/examples/introduction.rst b/docs/examples/introduction.rst
diff --git a/docs/examples/prediction/command_line.rst b/docs/examples/prediction/command_line.rst
diff --git a/docs/examples/prediction/python.rst b/docs/examples/prediction/python.rst
diff --git a/docs/examples/preprocessing/command_line.rst b/docs/examples/preprocessing/command_line.rst
diff --git a/docs/examples/preprocessing/python.rst b/docs/examples/preprocessing/python.rst
diff --git a/docs/examples/python/index.rst b/docs/examples/python/index.rst
diff --git a/docs/examples/python/lp_custom.rst b/docs/examples/python/lp_custom.rst
diff --git a/docs/examples/python/lp_fb15k237.rst b/docs/examples/python/lp_fb15k237.rst
diff --git a/docs/examples/python/nc_ogbn_arxiv.rst b/docs/examples/python/nc_ogbn_arxiv.rst
diff --git a/docs/export_and_inference/index.rst b/docs/export_and_inference/index.rst
diff --git a/docs/export_and_inference/marius_postprocess.rst b/docs/export_and_inference/marius_postprocess.rst
diff --git a/docs/export_and_inference/marius_predict.rst b/docs/export_and_inference/marius_predict.rst
diff --git a/docs/favicon.ico b/docs/favicon.ico
diff --git a/docs/graph_learning/decoders.rst b/docs/graph_learning/decoders.rst
diff --git a/docs/graph_learning/downstream_tasks.rst b/docs/graph_learning/downstream_tasks.rst
diff --git a/docs/graph_learning/encoders.rst b/docs/graph_learning/encoders.rst
diff --git a/docs/graph_learning/index.rst b/docs/graph_learning/index.rst
diff --git a/docs/graph_learning/intro.rst b/docs/graph_learning/intro.rst
diff --git a/docs/graph_learning/learning_tasks.rst b/docs/graph_learning/learning_tasks.rst
diff --git a/docs/index.rst b/docs/index.rst
diff --git a/docs/introduction.rst b/docs/introduction.rst
diff --git a/docs/marius_logo.png b/docs/marius_logo.png
diff --git a/docs/marius_logo.svg b/docs/marius_logo.svg
diff --git a/docs/marius_logo_scaled.png b/docs/marius_logo_scaled.png
diff --git a/docs/preprocess_datasets/built_in.rst b/docs/preprocess_datasets/built_in.rst
diff --git a/docs/preprocess_datasets/command_line.rst b/docs/preprocess_datasets/command_line.rst
diff --git a/docs/preprocess_datasets/index.rst b/docs/preprocess_datasets/index.rst
diff --git a/docs/preprocess_datasets/python.rst b/docs/preprocess_datasets/python.rst
diff --git a/docs/python_api/configuration/index.rst b/docs/python_api/configuration/index.rst
diff --git a/docs/python_api/data/batch.rst b/docs/python_api/data/batch.rst
diff --git a/docs/python_api/data/dataloader.rst b/docs/python_api/data/dataloader.rst
diff --git a/docs/python_api/data/dense_graph.rst b/docs/python_api/data/dense_graph.rst
diff --git a/docs/python_api/data/graph.rst b/docs/python_api/data/graph.rst
diff --git a/docs/python_api/data/index.rst b/docs/python_api/data/index.rst
diff --git a/docs/python_api/data/samplers/edge.rst b/docs/python_api/data/samplers/edge.rst
diff --git a/docs/python_api/data/samplers/index.rst b/docs/python_api/data/samplers/index.rst
diff --git a/docs/python_api/data/samplers/negative.rst b/docs/python_api/data/samplers/negative.rst
diff --git a/docs/python_api/data/samplers/neighbor.rst b/docs/python_api/data/samplers/neighbor.rst
diff --git a/docs/python_api/index.rst b/docs/python_api/index.rst
diff --git a/docs/python_api/manager/index.rst b/docs/python_api/manager/index.rst
diff --git a/docs/python_api/nn/activation.rst b/docs/python_api/nn/activation.rst
diff --git a/docs/python_api/nn/decoders/decoder.rst b/docs/python_api/nn/decoders/decoder.rst
diff --git a/docs/python_api/nn/decoders/edge/comparators.rst b/docs/python_api/nn/decoders/edge/comparators.rst
diff --git a/docs/python_api/nn/decoders/edge/complex.rst b/docs/python_api/nn/decoders/edge/complex.rst
diff --git a/docs/python_api/nn/decoders/edge/distmult.rst b/docs/python_api/nn/decoders/edge/distmult.rst
diff --git a/docs/python_api/nn/decoders/edge/edge_decoder.rst b/docs/python_api/nn/decoders/edge/edge_decoder.rst
diff --git a/docs/python_api/nn/decoders/edge/index.rst b/docs/python_api/nn/decoders/edge/index.rst
diff --git a/docs/python_api/nn/decoders/edge/relation_operators.rst b/docs/python_api/nn/decoders/edge/relation_operators.rst
diff --git a/docs/python_api/nn/decoders/edge/transe.rst b/docs/python_api/nn/decoders/edge/transe.rst
diff --git a/docs/python_api/nn/decoders/index.rst b/docs/python_api/nn/decoders/index.rst
diff --git a/docs/python_api/nn/decoders/node/index.rst b/docs/python_api/nn/decoders/node/index.rst
diff --git a/docs/python_api/nn/decoders/node/node_decoder.rst b/docs/python_api/nn/decoders/node/node_decoder.rst
diff --git a/docs/python_api/nn/decoders/node/noop_node_decoder.rst b/docs/python_api/nn/decoders/node/noop_node_decoder.rst
diff --git a/docs/python_api/nn/encoders/general_encoder.rst b/docs/python_api/nn/encoders/general_encoder.rst
diff --git a/docs/python_api/nn/encoders/index.rst b/docs/python_api/nn/encoders/index.rst
diff --git a/docs/python_api/nn/index.rst b/docs/python_api/nn/index.rst
diff --git a/docs/python_api/nn/initialization.rst b/docs/python_api/nn/initialization.rst
diff --git a/docs/python_api/nn/layers/embedding.rst b/docs/python_api/nn/layers/embedding.rst
diff --git a/docs/python_api/nn/layers/feature.rst b/docs/python_api/nn/layers/feature.rst
diff --git a/docs/python_api/nn/layers/gnn.rst b/docs/python_api/nn/layers/gnn.rst
diff --git a/docs/python_api/nn/layers/index.rst b/docs/python_api/nn/layers/index.rst
diff --git a/docs/python_api/nn/layers/layer.rst b/docs/python_api/nn/layers/layer.rst
diff --git a/docs/python_api/nn/layers/reduction.rst b/docs/python_api/nn/layers/reduction.rst
diff --git a/docs/python_api/nn/loss.rst b/docs/python_api/nn/loss.rst
diff --git a/docs/python_api/nn/model.rst b/docs/python_api/nn/model.rst
diff --git a/docs/python_api/nn/optim.rst b/docs/python_api/nn/optim.rst
diff --git a/docs/python_api/pipeline/evaluator.rst b/docs/python_api/pipeline/evaluator.rst
diff --git a/docs/python_api/pipeline/graph_encoder.rst b/docs/python_api/pipeline/graph_encoder.rst
diff --git a/docs/python_api/pipeline/index.rst b/docs/python_api/pipeline/index.rst
diff --git a/docs/python_api/pipeline/trainer.rst b/docs/python_api/pipeline/trainer.rst
diff --git a/docs/python_api/reporting/index.rst b/docs/python_api/reporting/index.rst
diff --git a/docs/python_api/reporting/metrics.rst b/docs/python_api/reporting/metrics.rst
diff --git a/docs/python_api/reporting/reporters.rst b/docs/python_api/reporting/reporters.rst
diff --git a/docs/python_api/storage/graph_storage.rst b/docs/python_api/storage/graph_storage.rst
diff --git a/docs/python_api/storage/index.rst b/docs/python_api/storage/index.rst
diff --git a/docs/python_api/storage/storage.rst b/docs/python_api/storage/storage.rst
diff --git a/docs/python_api/tools/configuration/constants.rst b/docs/python_api/tools/configuration/constants.rst
diff --git a/docs/python_api/tools/configuration/datatypes.rst b/docs/python_api/tools/configuration/datatypes.rst
diff --git a/docs/python_api/tools/configuration/index.rst b/docs/python_api/tools/configuration/index.rst
diff --git a/docs/python_api/tools/configuration/marius_config.rst b/docs/python_api/tools/configuration/marius_config.rst
diff --git a/docs/python_api/tools/index.rst b/docs/python_api/tools/index.rst
diff --git a/docs/python_api/tools/preprocess/converters/index.rst b/docs/python_api/tools/preprocess/converters/index.rst
diff --git a/docs/python_api/tools/preprocess/datasets/index.rst b/docs/python_api/tools/preprocess/datasets/index.rst
diff --git a/docs/python_api/tools/preprocess/index.rst b/docs/python_api/tools/preprocess/index.rst
diff --git a/docs/python_api/tools/preprocess/partitioners/index.rst b/docs/python_api/tools/preprocess/partitioners/index.rst
diff --git a/docs/python_api/tools/preprocess/readers/index.rst b/docs/python_api/tools/preprocess/readers/index.rst
diff --git a/docs/python_api/tools/preprocess/writers/index.rst b/docs/python_api/tools/preprocess/writers/index.rst
diff --git a/docs/quickstart.rst b/docs/quickstart.rst
diff --git a/examples/configuration/custom_lp.yaml b/examples/configuration/custom_lp.yaml
diff --git a/examples/configuration/custom_nc.yaml b/examples/configuration/custom_nc.yaml
diff --git a/examples/configuration/fb15k_237.yaml b/examples/configuration/fb15k_237.yaml
diff --git a/examples/configuration/ogbn_arxiv.yaml b/examples/configuration/ogbn_arxiv.yaml
diff --git a/examples/configuration/sakila.yaml b/examples/configuration/sakila.yaml
diff --git a/examples/db2graph/dockerfile b/examples/db2graph/dockerfile
diff --git a/examples/db2graph/run.sh b/examples/db2graph/run.sh
diff --git a/examples/docker/README.md b/examples/docker/README.md
diff --git a/examples/docker/cpu_ubuntu/dockerfile b/examples/docker/cpu_ubuntu/dockerfile
diff --git a/examples/docker/gpu_ubuntu/dockerfile b/examples/docker/gpu_ubuntu/dockerfile
diff --git a/examples/preprocessing/custom_dataset.py b/examples/preprocessing/custom_dataset.py
diff --git a/examples/python/custom.py b/examples/python/custom.py
diff --git a/examples/python/custom_lp.py b/examples/python/custom_lp.py
diff --git a/examples/python/custom_nc_graphsage.py b/examples/python/custom_nc_graphsage.py
diff --git a/examples/python/fb15k_237.py b/examples/python/fb15k_237.py
diff --git a/examples/python/fb15k_237_gpu.py b/examples/python/fb15k_237_gpu.py
diff --git a/examples/python/ogbn_arxiv_nc.py b/examples/python/ogbn_arxiv_nc.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/setup.cfg b/setup.cfg
diff --git a/setup.py b/setup.py
diff --git a/src/__init__.py b/src/__init__.py
diff --git a/src/cpp/cmake/FindSphinx.cmake b/src/cpp/cmake/FindSphinx.cmake
diff --git a/src/cpp/include/common/datatypes.h b/src/cpp/include/common/datatypes.h
diff --git a/src/cpp/include/common/exception.h b/src/cpp/include/common/exception.h
diff --git a/src/cpp/include/common/pybind_headers.h b/src/cpp/include/common/pybind_headers.h
diff --git a/src/cpp/include/common/util.h b/src/cpp/include/common/util.h
diff --git a/src/cpp/include/configuration/config.h b/src/cpp/include/configuration/config.h
@@ -143,6 +143,7 @@ struct TrainingConfig {
     int batch_size;
     shared_ptr<NegativeSamplingConfig> negative_sampling = nullptr;
     int num_epochs;
+    int batches_per_epoch; 
     shared_ptr<PipelineConfig> pipeline = nullptr;
     int epochs_per_shuffle;
     int logs_per_epoch;
@@ -157,6 +158,7 @@ struct EvaluationConfig {
     shared_ptr<NegativeSamplingConfig> negative_sampling = nullptr;
     shared_ptr<PipelineConfig> pipeline = nullptr;
     int epochs_per_eval;
+    int batches_per_epoch; 
     string checkpoint_dir;
     bool full_graph_evaluation;
 };

diff --git a/src/cpp/include/configuration/constants.h b/src/cpp/include/configuration/constants.h
diff --git a/src/cpp/include/configuration/options.h b/src/cpp/include/configuration/options.h
diff --git a/src/cpp/include/configuration/util.h b/src/cpp/include/configuration/util.h
diff --git a/src/cpp/include/data/batch.h b/src/cpp/include/data/batch.h
diff --git a/src/cpp/include/data/dataloader.h b/src/cpp/include/data/dataloader.h
diff --git a/src/cpp/include/data/graph.h b/src/cpp/include/data/graph.h
diff --git a/src/cpp/include/data/ordering.h b/src/cpp/include/data/ordering.h
diff --git a/src/cpp/include/data/samplers/edge.h b/src/cpp/include/data/samplers/edge.h
diff --git a/src/cpp/include/data/samplers/negative.h b/src/cpp/include/data/samplers/negative.h
diff --git a/src/cpp/include/data/samplers/neighbor.h b/src/cpp/include/data/samplers/neighbor.h
diff --git a/src/cpp/include/marius.h b/src/cpp/include/marius.h
diff --git a/src/cpp/include/nn/activation.h b/src/cpp/include/nn/activation.h
diff --git a/src/cpp/include/nn/decoders/decoder.h b/src/cpp/include/nn/decoders/decoder.h
diff --git a/src/cpp/include/nn/decoders/edge/comparators.h b/src/cpp/include/nn/decoders/edge/comparators.h
diff --git a/src/cpp/include/nn/decoders/edge/complex.h b/src/cpp/include/nn/decoders/edge/complex.h
diff --git a/src/cpp/include/nn/decoders/edge/decoder_methods.h b/src/cpp/include/nn/decoders/edge/decoder_methods.h
diff --git a/src/cpp/include/nn/decoders/edge/distmult.h b/src/cpp/include/nn/decoders/edge/distmult.h
diff --git a/src/cpp/include/nn/decoders/edge/edge_decoder.h b/src/cpp/include/nn/decoders/edge/edge_decoder.h
diff --git a/src/cpp/include/nn/decoders/edge/relation_operators.h b/src/cpp/include/nn/decoders/edge/relation_operators.h
diff --git a/src/cpp/include/nn/decoders/edge/transe.h b/src/cpp/include/nn/decoders/edge/transe.h
diff --git a/src/cpp/include/nn/decoders/node/node_decoder.h b/src/cpp/include/nn/decoders/node/node_decoder.h
diff --git a/src/cpp/include/nn/decoders/node/noop_node_decoder.h b/src/cpp/include/nn/decoders/node/noop_node_decoder.h
diff --git a/src/cpp/include/nn/encoders/encoder.h b/src/cpp/include/nn/encoders/encoder.h
diff --git a/src/cpp/include/nn/initialization.h b/src/cpp/include/nn/initialization.h
diff --git a/src/cpp/include/nn/layers/embedding/embedding.h b/src/cpp/include/nn/layers/embedding/embedding.h
diff --git a/src/cpp/include/nn/layers/feature/feature.h b/src/cpp/include/nn/layers/feature/feature.h
diff --git a/src/cpp/include/nn/layers/gnn/gat_layer.h b/src/cpp/include/nn/layers/gnn/gat_layer.h
diff --git a/src/cpp/include/nn/layers/gnn/gcn_layer.h b/src/cpp/include/nn/layers/gnn/gcn_layer.h
diff --git a/src/cpp/include/nn/layers/gnn/gnn_layer.h b/src/cpp/include/nn/layers/gnn/gnn_layer.h
diff --git a/src/cpp/include/nn/layers/gnn/graph_sage_layer.h b/src/cpp/include/nn/layers/gnn/graph_sage_layer.h
diff --git a/src/cpp/include/nn/layers/gnn/layer_helpers.h b/src/cpp/include/nn/layers/gnn/layer_helpers.h
diff --git a/src/cpp/include/nn/layers/gnn/rgcn_layer.h b/src/cpp/include/nn/layers/gnn/rgcn_layer.h
diff --git a/src/cpp/include/nn/layers/layer.h b/src/cpp/include/nn/layers/layer.h
diff --git a/src/cpp/include/nn/layers/reduction/concat.h b/src/cpp/include/nn/layers/reduction/concat.h
diff --git a/src/cpp/include/nn/layers/reduction/linear.h b/src/cpp/include/nn/layers/reduction/linear.h
diff --git a/src/cpp/include/nn/layers/reduction/reduction_layer.h b/src/cpp/include/nn/layers/reduction/reduction_layer.h
diff --git a/src/cpp/include/nn/loss.h b/src/cpp/include/nn/loss.h
diff --git a/src/cpp/include/nn/model.h b/src/cpp/include/nn/model.h
diff --git a/src/cpp/include/nn/model_helpers.h b/src/cpp/include/nn/model_helpers.h
diff --git a/src/cpp/include/nn/optim.h b/src/cpp/include/nn/optim.h
diff --git a/src/cpp/include/nn/regularizer.h b/src/cpp/include/nn/regularizer.h
diff --git a/src/cpp/include/pipeline/evaluator.h b/src/cpp/include/pipeline/evaluator.h
@@ -15,6 +15,9 @@
   The evaluator runs the evaluation process using the given model and dataset.
 */
 class Evaluator {
+   protected:
+    int batches_per_epoch_;
+
    public:
     shared_ptr<DataLoader> dataloader_;
 
@@ -31,7 +34,7 @@ class PipelineEvaluator : public Evaluator {
     shared_ptr<Pipeline> pipeline_;
 
    public:
-    PipelineEvaluator(shared_ptr<DataLoader> dataloader, shared_ptr<Model> model, shared_ptr<PipelineConfig> pipeline_config);
+    PipelineEvaluator(shared_ptr<DataLoader> dataloader, shared_ptr<Model> model, shared_ptr<PipelineConfig> pipeline_config, int batches_per_epoch = -1);
 
     void evaluate(bool validation) override;
 };
@@ -40,7 +43,7 @@ class SynchronousEvaluator : public Evaluator {
     shared_ptr<Model> model_;
 
    public:
-    SynchronousEvaluator(shared_ptr<DataLoader> dataloader, shared_ptr<Model> model);
+    SynchronousEvaluator(shared_ptr<DataLoader> dataloader, shared_ptr<Model> model, int batches_per_epoch = -1);
 
     void evaluate(bool validation) override;
 };

diff --git a/src/cpp/include/pipeline/graph_encoder.h b/src/cpp/include/pipeline/graph_encoder.h
diff --git a/src/cpp/include/pipeline/pipeline.h b/src/cpp/include/pipeline/pipeline.h
diff --git a/src/cpp/include/pipeline/pipeline_constants.h b/src/cpp/include/pipeline/pipeline_constants.h
diff --git a/src/cpp/include/pipeline/pipeline_cpu.h b/src/cpp/include/pipeline/pipeline_cpu.h
diff --git a/src/cpp/include/pipeline/pipeline_gpu.h b/src/cpp/include/pipeline/pipeline_gpu.h
diff --git a/src/cpp/include/pipeline/pipeline_monitor.h b/src/cpp/include/pipeline/pipeline_monitor.h
diff --git a/src/cpp/include/pipeline/queue.h b/src/cpp/include/pipeline/queue.h
diff --git a/src/cpp/include/pipeline/trainer.h b/src/cpp/include/pipeline/trainer.h
@@ -12,6 +12,9 @@
   The trainer runs the training process using the given model for the specified number of epochs.
 */
 class Trainer {
+   protected:
+    int batches_per_epoch_;
+
    public:
     shared_ptr<DataLoader> dataloader_;
     shared_ptr<ProgressReporter> progress_reporter_;
@@ -29,7 +32,7 @@ class PipelineTrainer : public Trainer {
     shared_ptr<Pipeline> pipeline_;
 
    public:
-    PipelineTrainer(shared_ptr<DataLoader> dataloader, std::shared_ptr<Model> model, shared_ptr<PipelineConfig> pipeline_config, int logs_per_epoch = 10);
+    PipelineTrainer(shared_ptr<DataLoader> dataloader, std::shared_ptr<Model> model, shared_ptr<PipelineConfig> pipeline_config, int logs_per_epoch = 10, int batches_per_epoch = -1);
 
     void train(int num_epochs = 1) override;
 };
@@ -38,7 +41,7 @@ class SynchronousTrainer : public Trainer {
     std::shared_ptr<Model> model_;
 
    public:
-    SynchronousTrainer(shared_ptr<DataLoader> dataloader, std::shared_ptr<Model> model, int logs_per_epoch = 10);
+    SynchronousTrainer(shared_ptr<DataLoader> dataloader, std::shared_ptr<Model> model, int logs_per_epoch = 10, int batches_per_epoch = -1);
 
     void train(int num_epochs = 1) override;
 };

diff --git a/src/cpp/include/reporting/logger.h b/src/cpp/include/reporting/logger.h
diff --git a/src/cpp/include/reporting/reporting.h b/src/cpp/include/reporting/reporting.h
diff --git a/src/cpp/include/storage/buffer.h b/src/cpp/include/storage/buffer.h
diff --git a/src/cpp/include/storage/checkpointer.h b/src/cpp/include/storage/checkpointer.h
diff --git a/src/cpp/include/storage/graph_storage.h b/src/cpp/include/storage/graph_storage.h
@@ -11,9 +11,11 @@
 
 struct GraphModelStoragePtrs {
     shared_ptr<Storage> edges = nullptr;
+    shared_ptr<Storage> edges_weights = nullptr;
     shared_ptr<Storage> train_edges = nullptr;
-    shared_ptr<Storage> train_edges_dst_sort = nullptr;
     shared_ptr<Storage> train_edges_weights = nullptr;
+    shared_ptr<Storage> train_edges_dst_sort = nullptr;
+    shared_ptr<Storage> train_edges_dst_sort_weights = nullptr;
     shared_ptr<Storage> validation_edges = nullptr;
     shared_ptr<Storage> validation_edges_weights = nullptr;
     shared_ptr<Storage> test_edges = nullptr;

diff --git a/src/cpp/include/storage/io.h b/src/cpp/include/storage/io.h
diff --git a/src/cpp/include/storage/storage.h b/src/cpp/include/storage/storage.h
diff --git a/src/cpp/python_bindings/configuration/config_wrap.cpp b/src/cpp/python_bindings/configuration/config_wrap.cpp
diff --git a/src/cpp/python_bindings/configuration/options_wrap.cpp b/src/cpp/python_bindings/configuration/options_wrap.cpp
diff --git a/src/cpp/python_bindings/configuration/wrap.cpp b/src/cpp/python_bindings/configuration/wrap.cpp
diff --git a/src/cpp/python_bindings/data/batch_wrap.cpp b/src/cpp/python_bindings/data/batch_wrap.cpp
Original file line number	Diff line number	Diff line change
Expand Up		@@ -171,3 +171,5 @@ Thumbs.db


		# End of https://www.toptal.com/developers/gitignore/api/python

		test_datasets