How facebookresearch/neural-rewriter works
1. What this repository is
- This is a PyTorch-based research codebase for the “Neural Rewriter” framework from the NeurIPS 2019 paper.
- It implements one shared pattern: take a combinatorial instance, encode its current state, choose a local rewrite/action, and iteratively improve the solution.
2. Top-level task split
- The repo is organized around three problem domains:
- Halide expression simplification (run_Halide.py, Halide_search.py)
- Job scheduling (run_jsp.py, jsp_nonNN_baselines.py)
- Vehicle routing (run_vrp.py)
- Each task has its own model, input encoder, and rewriter, but they share the same training/evaluation scaffolding.
3. Shared model architecture
- src/models/BaseModel.py defines the common model interface.
- Task-specific models extend it:
- HalideModel.py
- jspModel.py
- vrpModel.py
- The shared pattern is: parse instance → encode state → score candidate rewrites/actions → select and apply a rewrite.
4. Input representation layer
- src/models/data_utils/ is the preprocessing boundary between raw datasets and neural models.
- It includes different structured representations:
- Tree.py for expression trees
- Dag.py for DAG-like structures
- Seq.py for sequential representations
- parser.py, data_utils.py, utils.py for loading and transforming examples
- This layer is important because the three tasks do not share the same raw state format.
5. Task-specific encoders
- src/models/modules/ contains the neural feature extractors:
- HalideInputEncoder.py
- jspInputEncoder.py
- vrpInputEncoder.py
- mlp.py as a generic building block
- These encoders convert structured problem states into fixed-size embeddings used by the policy/model head.
6. Rewrite/action selection
- src/models/rewriter/ holds the logic that maps model outputs to actual local edits:
- HalideRewriter.py
- jspRewriter.py
- vrpRewriter.py
- This is the key architectural boundary: the neural network does not directly solve the whole problem, it proposes local transformations that the task-specific rewriter applies.
7. Training and experiment control
- src/models/model_utils/ contains shared experiment infrastructure:
- supervisor.py for training/evaluation orchestration
- logger.py for logging
- src/arguments.py centralizes CLI options such as dataset paths, model directory, evaluation mode, epochs, and checkpointing cadence.
- The run scripts wire together arguments, data loading, model construction, and supervisor logic.
8. Halide pipeline
- run_Halide.py is the main training/evaluation entry point for expression simplification.
- Halide_search.py provides a non-neural search baseline over Halide rewrite rules.
- This task depends on expression-tree parsing and on the Halide rule set conceptually, but the repo’s code focuses on learning and search around those rewrite traces.
9. Job scheduling pipeline
- run_jsp.py trains/evaluates the neural rewriter for online scheduling with a bounded pending queue.
- jsp_nonNN_baselines.py runs heuristic baselines such as SJF, EJF, and random.
- The data flow here is: generated scheduling trace → structured state → encoder → rewrite/scheduling decision → updated schedule state.
10. Vehicle routing pipeline
- run_vrp.py is the entry point for routing experiments.
- The vehicle-routing model works over customer/depot state and capacity constraints, using the VRP-specific encoder and rewriter.
- External baselines mentioned in the README include OR-Tools and reinforcement-learning solvers, but the repo’s core path is the neural rewrite loop.
11. Dataset generation
- data_generator/jspDatagen.py and data_generator/vrpDatagen.py synthesize training/evaluation data for the scheduling and routing tasks.
- These generators are part of the offline data pipeline, not the model runtime.
- For Halide, the dataset comes from generated expressions and rewrite traces derived from the Halide ecosystem.
12. External dependencies and boundaries
- PyTorch is the core runtime for all neural components.
- Halide, Z3, and OR-Tools appear as external systems for dataset creation or baseline comparisons, not as internal subsystems.
- The architectural boundary is clear: external solvers/heuristics provide data or reference performance, while the repo’s own code implements the neural rewrite policy and its task-specific application layer.
Main entry points: supervisor.py Supervisor · data_utils.py Data utils · BaseModel.py Base model · arguments.py CLI args
Components 26
Shared core
Common model, data, and training infrastructure used across tasks.
- CLI args · Central command-line options for training, evaluation, and checkpoints.
src/arguments.py - Base model · Common model abstraction extended by task-specific models.
src/models/BaseModel.py - Supervisor · Shared experiment loop for training, evaluation, and checkpoint handling.
src/models/model_utils/supervisor.py - Logger · Shared experiment logging utilities.
src/models/model_utils/logger.py - Data utils · Boundary layer for loading, parsing, and transforming raw problem instances.
src/models/data_utils/data_utils.py - Parser · Parses task-specific raw inputs into structured representations.
src/models/data_utils/parser.py - Tree · Structured tree representation for Halide expressions.
src/models/data_utils/Tree.py - DAG · DAG-like state representation used by the data layer.
src/models/data_utils/Dag.py - Sequence · Sequential state representation used by the data layer.
src/models/data_utils/Seq.py - MLP · Generic feed-forward building block used inside task models.
src/models/modules/mlp.py
Halide
Expression simplification pipeline and rewrite search.
- Halide model · Neural model for expression simplification.
src/models/HalideModel.py - Halide encoder · Encodes Halide expression state into model features.
src/models/modules/HalideInputEncoder.py - Halide rewriter · Maps model scores to concrete Halide rewrite actions.
src/models/rewriter/HalideRewriter.py - Halide run · Training and evaluation entry point for Halide experiments.
src/run_Halide.py - Halide search · Non-neural search baseline over Halide rewrite rules.
src/Halide_search.py
Job scheduling
Scheduling pipeline, model, and heuristics.
- JSP model · Neural model for job scheduling decisions.
src/models/jspModel.py - JSP encoder · Encodes scheduling state and pending jobs for the policy.
src/models/modules/jspInputEncoder.py - JSP rewriter · Selects and applies scheduling actions to the current state.
src/models/rewriter/jspRewriter.py - JSP run · Training and evaluation entry point for job scheduling.
src/run_jsp.py - JSP baselines · Non-neural scheduling baselines such as SJF, EJF, and random.
src/jsp_nonNN_baselines.py
Vehicle routing
Routing pipeline and model.
- VRP model · Neural model for vehicle routing decisions.
src/models/vrpModel.py - VRP encoder · Encodes routing state, customers, depot, and capacity features.
src/models/modules/vrpInputEncoder.py - VRP rewriter · Selects route construction or local routing actions.
src/models/rewriter/vrpRewriter.py - VRP run · Training and evaluation entry point for routing experiments.
src/run_vrp.py
Data generation
Offline synthetic data generators for training and evaluation.
- JSP data gen · Offline generator for scheduling traces and training data.
data_generator/jspDatagen.py - VRP data gen · Offline generator for routing instances and training data.
data_generator/vrpDatagen.py
Connections 30
- CLI args → Halide run: configures · CLI options flow into the Halide entry point.
- CLI args → JSP run: configures · CLI options flow into the JSP entry point.
- CLI args → VRP run: configures · CLI options flow into the VRP entry point.
- Halide run → Supervisor: drives · Halide experiments are executed through the shared supervisor.
- JSP run → Supervisor: drives · JSP experiments are executed through the shared supervisor.
- VRP run → Supervisor: drives · VRP experiments are executed through the shared supervisor.
- Supervisor → Logger: logs to · Training and evaluation metrics are recorded by the shared logger.
- Supervisor → Base model: trains · Supervisor operates on the common model interface.
- Data utils → Parser: parses · Raw inputs are turned into structured instances.
- Data utils → Tree: represents · Expression trees are one of the core structured forms.
- Data utils → DAG: represents · Graph-like state is normalized through DAG representations.
- Data utils → Sequence: represents · Sequential state is normalized through sequence representations.
- Halide model → Base model: extends · Halide-specific model implements the shared interface.
- Halide model → Halide encoder: uses · Halide model consumes encoded expression state.
- Halide model → Halide rewriter: produces actions · Model outputs are turned into rewrite operations.
- Halide encoder → MLP: builds on · Encoder uses generic neural blocks where needed.
- Halide rewriter → Tree: applies to · Rewrite actions operate on expression trees.
- Halide search → Tree: searches · Baseline search explores expression rewrite traces.
- JSP model → Base model: extends · JSP-specific model implements the shared interface.
- JSP model → JSP encoder: uses · JSP model consumes scheduling-state embeddings.
- JSP model → JSP rewriter: produces actions · Model outputs are mapped to scheduling decisions.
- JSP encoder → MLP: builds on · Encoder uses generic neural blocks where needed.
- JSP rewriter → Sequence: applies to · Scheduling actions update sequential state.
- VRP model → Base model: extends · VRP-specific model implements the shared interface.
- VRP model → VRP encoder: uses · VRP model consumes routing-state embeddings.
- VRP model → VRP rewriter: produces actions · Model outputs are mapped to routing decisions.
- VRP encoder → MLP: builds on · Encoder uses generic neural blocks where needed.
- VRP rewriter → DAG: applies to · Routing actions update graph-structured route state.
- JSP data gen → JSP run: feeds · Synthetic scheduling traces supply JSP experiments.
- VRP data gen → VRP run: feeds · Synthetic routing instances supply VRP experiments.