InhwanBae / ENet-SAD_Pytorch

Pytorch implementation of "Learning Lightweight Lane Detection CNNs by Self Attention Distillation (ICCV 2019)"
MIT License
76 stars 18 forks source link
deep-learning enet lane-detection python pytorch pytorch-implementation self-attention-distillation semantic-segmentation

ENet-SAD Pytorch

Pytorch implementation of "Learning Lightweight Lane Detection CNNs by Self Attention Distillation (ICCV 2019)"

drawing

News

You can find the previous version here

Demo

Video

demo_gif

Demo trained with CULane dataset & tested with \driver_193_90frame\06051123_0635.MP4

gpu_runtime: 0.016253232955932617 FPS: 61

total_runtime: 0.017553091049194336 FPS: 56 on RTX 2080 TI

Comparison

Category 40k episode (before SAD) 60k episode (after SAD)
Image img1 img2
Lane img3 img4

Train

Requirements

Datasets

You need to change the correct dataset path in ./config.py

Dataset_Path = dict(
    CULane = "/workspace/CULANE_DATASET",
    Tusimple = "/workspace/TUSIMPLE_DATASET",
    BDD100K = "/workspace/BDD100K_DATASET",
    VPGNet = "/workspace/VPGNet_DATASET"
)

Training

First, change some hyperparameters in ./experiments/*/cfg.json

{
  "model": "enet_sad",               <- "scnn" or "enet_sad"
  "dataset": {
    "dataset_name": "CULane",        <- "CULane" or "Tusimple"
    "batch_size": 12,
    "resize_shape": [800, 288]       <- [800, 288] with CULane, [640, 368] with Tusimple, and [640, 360] with BDD100K
                                        This size is defined in the ENet-SAD paper, any size is fine if it is a multiple of 8.
  },
  ...
}

And then, start training with train.py

python train.py --exp_dir ./experiments/exp1
Optional

If you write your own code using ENet-SAD model, you can init ENet_SAD with some parameters.

class ENet_SAD(nn.Module):
    """
    Generate the ENet-SAD model.
    Keyword arguments:
    - input_size (tuple): Size of the input image to segment.
    - encoder_relu (bool, optional): When ``True`` ReLU is used as the
    activation function in the encoder blocks/layers; otherwise, PReLU
    is used. Default: False.
    - decoder_relu (bool, optional): When ``True`` ReLU is used as the
    activation function in the decoder blocks/layers; otherwise, PReLU
    is used. Default: True.
    - sad (bool, optional): When ``True``, SAD is added to model;
    otherwise, SAD is removed. Default: False.
    - weight_share (bool, optional): When ``True``, weights are shared
    with encoder blocks (E2, E3 and E4); otherwise, encoders will each
    learn with their own weights. Default: True.
    """

Performance

Will continue to be updated.

Acknowledgement

This repo is built upon official implementation ENet-SAD and based on PyTorch-ENet, SCNN_Pytorch.