A Pytorch Re-Implementation of TNT & VectorNet on Custom dataset with TensorRT deploy

Paper:

TNT: Target-driveN Trajectory Prediction

VectorNet: Encoding HD Maps and Agent Dynamics from Vectorized Representation

Thanks to this TNT implemention

I. Features

Contain follow features:

• my custom data feature preprocessor
• remove torch-geometric requirements
• batchify the data and compute subgraph in pipeline
• better visualize the evaluation result
• use tensorboard to visualize the loss & metric
• simplify the inference pipeline for deploy
• add output log during training
• deploy through TensorRT and libtorch

II. Train & Test

1. Add python path

export PYTHONPATH="${PYTHONPATH}:$(pwd)"

2. Train

python tools/train_vectornet.py -d ./mini_data -b 128 --lr 0.005
or
python tools/train_tnt.py -d ./mini_data -b 128 --lr 0.005

3. Test

python tools/test_vectornet.py -d ./mini_data -b 128 -rm weights/best_vectornet.pth
or
python tools/test_tnt.py -d ./mini_data -b 128 -rm weights/best_tnt.pth

III. Plot

ATTENTION: My custom dataset does not have HDMap, so I only use the perception lane for input. This is only for experiment !!!

IV.Tensorboard

V. Metric

On my private dataset, the metrics are:

{
    'minADE': 0.6467818400568369,
    'minFDE': 1.3008982805719813,
    'MR': 0.20051413881748073,
    'p-minADE': 2.153652389728194,
    'p-minFDE': 2.807768830243343,
    'p-MR': 0.7725899742930608,
    'brier-minADE': 1.1959656446840874,
    'brier-minFDE': 1.8500820851992275
}

VI. Deploy

1. Compute graph

Organize the inference computing pipeline as follows:

2. TRT deploy

Details can be found in tensorrt_deploy directory.

• Speed Test

VectorNet speed:

Libtorch are ABANDONED, please refer to TRT deploy!!!

3. Libtorch deploy

Only test on libtorch-cpu(version must > 1.12.0), average 0.4 ms/forward.

a. Export jit traced module.

python assets/tools/export/vectornet_export_v2.py

b. Compile cpp code.

cd assets/tools/export/cpp_libtorch > mkdir build && cd build > cmake .. > make

c. Excute

./vectornet ../../models/traced_vectornet.pt