royalskye / routing-mvmoe

[ICML 2024] "MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-Experts"

MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-Experts

                

The PyTorch Implementation of ICML 2024 -- MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-Experts. MVMoE is a unified neural solver that can cope with 16 VRP variants simultaneously, even in a zero-shot manner. Concretely, the training tasks include CVRP, OVRP, VRPB, VRPL, VRPTW, and OVRPTW. The test tasks include OVRPB, OVRPL, VRPBL, VRPBTW, VRPLTW, OVRPBL, OVRPBTW, OVRPLTW, VRPBLTW, and OVRPBLTW.

• 🇦🇹 We will attend ICML 2024. Welcome to stop by our poster (Session 6 @ Hall C 4-9 #1003) for discussion.
• 🚀 RL4CO adds support for MVMoE - example.
• 🎓 We release an awesome paper list on Foundation Model for Combinatorial Optimization - FM4CO.

Overview

Poster

Dependencies

• Python >= 3.8
• Pytorch >= 1.12

How to Run

Train

# Default: --problem_size=100 --pomo_size=100 --gpu_id=0
# 0. POMO
python train.py --problem={PROBLEM} --model_type=SINGLE

# 1. POMO-MTL
python train.py --problem=Train_ALL --model_type=MTL

# 2. MVMoE/4E 
python train.py --problem=Train_ALL --model_type=MOE --num_experts=4 --routing_level=node --routing_method=input_choice

# 3. MVMoE/4E-L
python train.py --problem=Train_ALL --model_type=MOE_LIGHT --num_experts=4 --routing_level=node --routing_method=input_choice

Evaluation

# 0. POMO
python test.py --problem={PROBLEM} --model_type=SINGLE --checkpoint={MODEL_PATH}

# 1. POMO-MTL
python test.py --problem=ALL --model_type=MTL --checkpoint={MODEL_PATH}

# 2. MVMoE/4E
python test.py --problem=ALL --model_type=MOE --num_experts=4 --routing_level=node --routing_method=input_choice --checkpoint={MODEL_PATH}

# 3. MVMoE/4E-L
python test.py --problem=ALL --model_type=MOE_LIGHT --num_experts=4 --routing_level=node --routing_method=input_choice --checkpoint={MODEL_PATH}

# 4. Evaluation on CVRPLIB
python test.py --problem=CVRP --model_type={MODEL_TYPE} --checkpoint={MODEL_PATH} --test_set_path=../data/CVRP-LIB

Baseline

# 0. LKH3 - Support for ["CVRP", "OVRP", "VRPL", "VRPTW"]
python LKH_baseline.py --problem={PROBLEM} --datasets={DATASET_PATH} -n=1000 --cpus=32 -runs=1 -max_trials=10000

# 1. HGS - Support for ["CVRP", "VRPTW"]
python HGS_baseline.py --problem={PROBLEM} --datasets={DATASET_PATH} -n=1000 --cpus=32 -max_iteration=20000

# 2. OR-Tools - Support for all 16 VRP variants
python OR-Tools_baseline.py --problem={PROBLEM} --datasets={DATASET_PATH} -n=1000 --cpus=32 -timelimit=20

How to Customize MoE

MoEs can be easily used in Transformer-based models by replacing a Linear/MLP with an MoE layer. The input and output dimensions are kept the same as the original layer. Below, we provide two examples of how to customize MoEs.

# 0. Our implementation based on https://github.com/davidmrau/mixture-of-experts
# Supported routing levels: node/instance/problem
# Supported routing methods: input_choice/expert_choice/soft_moe/random (only for node/instance gating levels)
from MOELayer import MoE
moe_layer = MoE(input_size={INPUT_DIM}, output_size={OUTPUT_DIM}, hidden_size={HIDDEN_DIM},
                num_experts={NUM_EXPERTS}, k=2, T=1.0, noisy_gating=True, 
                routing_level="node", routing_method="input_choice", moe_model="MLP")

# 1. tutel - https://github.com/microsoft/tutel
from tutel import moe as tutel_moe
moe_layer = tutel_moe.moe_layer(
                gate_type={'type': 'top', 'k': 2},
                model_dim={INPUT_DIM},
                experts={'type': 'ffn', 'count_per_node': {NUM_EXPERTS},
                         'hidden_size_per_expert': {HIDDEN_DIM},
                         'activation_fn': lambda x: torch.nn.functional.relu(x)},
            )

Citation

@inproceedings{zhou2024mvmoe,
title       ={MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-Experts},
author      ={Jianan Zhou and Zhiguang Cao and Yaoxin Wu and Wen Song and Yining Ma and Jie Zhang and Chi Xu},
booktitle   ={International Conference on Machine Learning},
year        ={2024}
}

Acknowledgments

• ICML 2024 Review
• https://github.com/yd-kwon/POMO
• https://github.com/davidmrau/mixture-of-experts

routing-mvmoe's Issues

About Checkpoints

Thanks for sharing checkpoints and codes!

However, checkpoints for POMO-CVRP can't be loaded into any models since mismatch between tensors.

In MLTEncoder, to encode location/demand information, codes will use nn.Embedding(5,128). However, in the checkpoint POMO-CVRP.pt, the shape is nn.Embedding(3,128). So there is a mismatch.

About Pre-trained Models

Thanks for sharing the pre-trained models.

But in the pretrained directory, I only saw two POMO based neural solvers: "POMO-CVRP" and "POMO-VRPTW". In the paper, there should be another four single models: POMO-OVRP, POMO-VRPL, POMO-VRPB, POMO-OVRPTW. I really appreciate if you could release them as well.