This repository provides the PyTorch implementation of the paper EurNet: Efficient Multi-Range Relational Modeling of Spatial Multi-Relational Data. This repository contains complete source code and model weights for the image classification experiments in the paper. Currently, we release all source code and model weights of EurNet for image modeling. The codes and model weights for protein and knowledge graph modeling will be released soon after refactor.
EurNet employs the Gated Relational Message Passing (GRMP) layer as its basic component, as graphically shown above. EurNet can be applied to various domains (e.g., images, protein structures and knowledge graphs) for efficient multi-relational modeling at scale. Here are the links to other applied domains/tasks of this project:
- EurNet for Image Classification (PyTorch Geometric implementation)
- EurNet for Object Detection
- EurNet for Semantic Segmentation
- EurNet for Protein Structure Modeling (Working, will release soon)
- EurNet for Knowledge Graph Reasoning (Working, will release soon)
- [2022/11/24] The initial release! We release all source code and model weights of EurNet for image classification (both TorchDrug and PyTorch Geometric implementations), object detection and semantic segmentation.
- Release code and model weights of EurNet for protein structure modeling.
- Release code and model weights of EurNet for knowledge graph reasoning.
Model | Training | #Params. (M) | FLOPs (G) | IN-1K Top-1 (%) | Config | Ckpt | Log |
---|---|---|---|---|---|---|---|
EurNet-T | 1K-scratch | 29 | 4.6 | 82.3 | config | ckpt | log |
EurNet-S | 1K-scratch | 50 | 8.8 | 83.6 | config | ckpt | log |
EurNet-B | 1K-scratch | 89 | 15.6 | 84.1 | config | ckpt | log |
EurNet-B | 22K-pretrain | 89 | 15.6 | - | config | ckpt | - |
EurNet-B/384 | 224to384-tune | 90 | 46.6 | 85.3 | config | ckpt | log |
EurNet-B | 22Kto1K-tune | 89 | 15.6 | 85.7 | config | ckpt | log |
EurNet-B/384 | 22Kto1K-tune | 90 | 46.6 | 87.0 | config | ckpt | log |
This repository is officially tested with the following environments:
- Linux
- Python 3.6+
- PyTorch 1.10.0
- CUDA 11.3
The environment could be prepared in the following steps:
- Create a virtual environment with conda:
conda create -n eurnet python=3.7.3 -y
conda activate eurnet
- Install PyTorch with the official instructions. For example:
conda install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge
- Install dependencies:
pip install -r requirements.txt
We support ImageNet ILSVRC 2012 for training from scratch or fine-tuning.
We recommend symlink the dataset folder to ./datasets/ImageNet1K
. The folder structure would be:
datasets/
ImageNet1K/
ILSVRC/
Annotations/
Data/
train/
val/
ImageSets/
meta/
After downloading and unzip the dataset, go to path ./datasets/ImageNet1K/ILSVRC/Data/val/
and move images to labeled sub-folders with this script.
We also support ImageNet22K for pre-training. We recommend symlink the dataset folder to ./datasets/imagenet22k/
. Move all images to the labeled sub-folders in the data path. Then, download the train-val split files (ILSVRC2011fall_whole_map_train.txt
& ILSVRC2011fall_whole_map_val.txt) and move to the data path.
The folder structure should be:
datasets/
imagenet22k/
ILSVRC2011fall_whole_map_train.txt
ILSVRC2011fall_whole_map_val.txt
n00004475/
n00005787
...
We provide an easy yaml based configuration file system. The config could be modified by command line arguments.
To run an experiment:
python3 launch.py --launch ./tools/train.py -c [config file] [config options]
The config options are in "key=value" format. For example, ouput_dir=your_path
and batch_size=64
.
Sub module is seperated by .
. For example, optimizer.name=AdamW
modifies the sub key name
in
optimizer
with value AdamW
.
A full example of training EurNet-T on ImageNet1K for 300 epochs:
python3 launch.py --launch ./tools/train.py -c configs/classification/eurnet_tiny_1k_300eps.yaml \
output_dir=./experiments/imagenet1k/EurNet-T/
Note that, all the training configuration files are in ./configs/classification/
. To reproduce the ImageNet classification results, please follow the corresponding config file.
launch.py
would automatically find a free port to launch single node experiments.
If multiple nodes training is necessary, the number of nodes --nn
, node rank --nr
, master port --port
and master address -ma
should be set.
Take training EurNet-T on ImageNet-1K with two nodes as an example, use the following commands at node1 and node2, respectively.
# use this command at node 1
python3 launch.py --nn 2 --nr 0 --port [port] -ma [address of node 0] --launch ./tools/train.py \
-c configs/classification/eurnet_tiny_1k_300eps.yaml
# use this command at node 2
python3 launch.py --nn 2 --nr 1 --port [port] -ma [address of node 0] --launch ./tools/train.py \
-c configs/classification/eurnet_tiny_1k_300eps.yaml
This repository is released under the MIT license as in the LICENSE file.
If you find this repository useful in your research, please cite the following paper:
@article{xu2022eurnet,
title={EurNet: Efficient Multi-Range Relational Modeling of Spatial Multi-Relational Data},
author={Xu, Minghao and Guo, Yuanfan and Xu, Yi and Tang, Jian and Chen, Xinlei and Tian, Yuandong},
journal={arXiv preprint arXiv:2211.12941},
year={2022}
}
The development of this project is guided by the following open-source projects. We would like to thank the authors for releasing the source code.