• Conda virtual environment or docker is required to setup the environment

Using conda environment

git clone https://github.com/j-marple-dev/AYolov2.git
cd AYolov2
./run_check.sh init
# Equivalent to
# conda env create -f environment.yml
# pre-commit install --hook-type pre-commit --hook-type pre-push

Using docker

Building a docker image

./run_docker.sh build
# You can add build options
# ./run_docker.sh build --no-cache

Running the container

This will mount current repository directory from local disk to docker image

./run_docker.sh run
# You can add running options
# ./run_docker.sh run -v $DATASET_PATH:/home/user/dataset

Executing the last running container

./run_docker.sh exec

• Example

  python3 train.py --model $MODEL_CONFIG_PATH --data $DATA_CONFIG_PATH --cfg $TRAIN_CONFIG_PATH
  # i.e.
  # python3 train.py --model res/configs/model/yolov5s.yaml --data res/configs/data/coco.yaml --cfg res/configs/cfg/train_config.yaml
  # Logging and upload trained weights to W&B
  # python3 train.py --model res/configs/model/yolov5s.yaml --wlog --wlog_name yolov5s

  Prepare dataset

  • Dataset config file

  train_path: "DATASET_ROOT/images/train"
  val_path: "DATASET_ROOT/images/val"
  
  # Classes
  nc: 10  # number of classes
  dataset: "DATASET_NAME"
  names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light']  # class names

  • Dataset directory structure
    • One of labels or segments directory must exist.
    • Training label type(labels or segments) will be specified in the training config.
    • images and labels or segments must have a matching filename with .txt extension.

  DATASET_ROOT
  │
  ├── images
  │   ├── train
  │   └── val
  ├── labels
  │   ├── train
  │   └── val
  ├── segments
  │   ├── train
  │   └── val

  Training config

  • Default training configurations are defined in train_config.yaml.
  • You may want to change batch_size, epochs, device, workers, label_type along with your model, dataset, and training hardware.
  • Be cautious to change other parameters. It may affect training results.
  Model config

  • Model is defined by yaml file with kindle
  • Please refer to https://github.com/JeiKeiLim/kindle
  Multi-GPU training

  • Please use torch.distributed.run module for multi-GPU Training

  python3 -m torch.distributed.run --nproc_per_node $N_GPU train.py --model $MODEL_CONFIG_PATH --data $DATA_CONFIG_PATH --cfg $TRAIN_CONFIG_PATH

  - N_GPU: Number of GPU to use
  

• Validate from local weights

python3 val.py --weights $WEIGHT_PATH --data-cfg $DATA_CONFIG_PATH

• You can pass W&B path to the weights argument.

python3 val.py --weights j-marple/AYolov2/179awdd1 --data-cfg $DATA_CONFIG_PATH

• TTA (Test Time Augmentation)
  • Simply pass --tta argument with --tta-cfg path
  • Default TTA configs are located in res/configs/cfg/tta.yaml

python3 val.py --weights $WEIGHT_PATH --data-cfg $DATA_CONFIG_PATH --tta --tta-cfg $TTA_CFG_PATH

• Validate with pycocotools (Only for COCO val2017 images) Future work: The val.py and val2.py should be merged together.

python3 val2.py --weights $WEIGHT_PATH --data $VAL_IMAGE_PATH --json-path $JSON_FILE_PATH

• You can export a trained model to TorchScript, ONNX, or TensorRT

• INT8 quantization is currently not supported (coming soon).

• Usage

python3 export.py --weights $WEIGHT_PATH --type [torchscript, ts, onnx, tensorrt, trt] --dtype [fp32, fp16, int8]

• Above command will generate both model and model config file.

  • Example) FP16, Batch size 8, Image size 640x640, TensorRT
    • model_fp16_8_640_640.trt
    • model_fp16_8_640_640_trt.yaml

    batch_size: 8
    conf_t: 0.001  # NMS confidence threshold
    dst: exp/  # Model location
    dtype: fp16  # Data type
    gpu_mem: 6  # GPU memory restriction
    img_height: 640
    img_width: 640
    iou_t: 0.65  # NMS IoU threshold
    keep_top_k: 100  # NMS top k parameter
    model_cfg: res/configs/model/yolov5x.yaml  # Base model config location
    opset: 11  # ONNX opset version
    rect: false  # Rectangular inference mode
    stride_size: 32  # Model stride size
    top_k: 512  # Pre-NMS top k parameter
    type: trt  # Model type
    verbose: 1  # Verbosity level
    weights: ./exp/yolov5x.pt  # Base model weight file location

• Once, model has been exported, you can run val.py with the exported model.

  • ONNX inference is currently not supported.

  python3 val.py --weights model_fp16_8_640_640.trt --data-cfg $DATA_CONFIG_PATH

• A trained model can be compressed via tensor decomposition.

• Decomposed conv is composed of 3 convolutions from 1 large convolution.

  • Example)
    • Original conv: 64x128x3x3
    • Decomposed conv: 64x32x1x1 -> 32x16x3x3 -> 16x128x1x1

• Usage

  python3 decompose_model.py --weights $WEIGHT_PATH --loss-thr $DECOMPOSE_LOSS_THRESHOLD --prune-step $PRUNING_STEP --data-cfg $DATA_CONFIG_PATH

  ...
  [  Original] # param: 86,749,405, mAP0.5: 0.678784398716757, Speed(pre-process, inference, NMS): 0.030, 21.180, 4.223
  [Decomposed] # param: 49,508,630, mAP0.5: 0.6707606125947304, Speed(pre-process, inference, NMS): 0.030, 20.274, 4.345
  Decomposition config saved to exp/decompose/val/2021_0000_runs/args.yaml
  Decomposed model saved to exp/decompose/val/2021_0000_runs/yolov5x_decomposed.pt

  • Passing prune-step to 0 will skip pruning optimization.

Summary of tensor decomposition process

1. Pass random tensor x to original conv (ŷ) and decomposed conv (ỹ)
2. Compute E = Error(ŷ, ỹ)
3. If E < loss-thr, use decomposed conv
4. Apply pruning ratio with binary search
5. Jump to 1 until differential of pruning ratio is less than prune-step

:: Note :: Decomposition process uses CPU only.

• An ad-hoc implementation of the knowledge distillation motivated from the method in "End-to-end semi-supervised object dection with soft teacher".
• Create pseudo-labels for "unlabeled dataset" using the inference of the "teacher" model.
• :: Note ::
  • Implemented to use the same dataset for the "training dataset" and the "unlabeled dataset". To use different datasets, the creation of the dataloader instance unlabeled_loader in distillation.py should be modified.
  • Teacher model weights are fixed during training student model. (In the original paper, teacher model is updated using "exponential moving averaging" the student model.)
• Usage

  python distillation.py --model res/configs/model/yolov5s.yaml \
                         --cfg res/configs/cfg/distillation.yam \
                         --data res/configs/data/coco.yaml \
                         --teacher {wandb_runpath_of_pretrained_model}

• Representations of a model can be automatically discovered from raw data by representation learning.

• You can apply SimpleRL or SimCLR to find better representations of the model with --rl-type option.

  • SimpleRL is a method to minimize a difference between last two representations of a model with L1 loss.
  • SimCLR is a simple framework for contrastive self-supervised learning of visual representations without requiring specialized architectures.

• Usage (default: base)

  • SimpleRL

    python train_repr.py --model res/configs/model/yolov5s_repr.yaml \
                         --data res/configs/data/coco_repr.yaml \
                         --cfg res/configs/cfg/train_config_repr.yaml \
                         --rl-type base

  • SimCLR

    python train_repr.py --model res/configs/model/simclr.yaml \
                         --data res/configs/data/coco_repr.yaml \
                         --cfg res/configs/cfg/train_config_simclr.yaml \
                         --rl-type simclr

If want to optimize NMS parameters(IoU threshold and confidence threshold), there are two ways to optimize.

Why there are two ways?

• There is an issue with YOLOv5 validation.
• It's ok with training or validating but the validation results are little different.

1. Optimize parameters with YOLOv5 validation.
2. Optimize parameters with COCO validation (pycocotools).

1. Optimize parameters with YOLOv5 validation.

python3 val_optimizer.py --weights ${WEIGHT_PATH | WANDB_PATH} --data-cfg $DATA_CONFIG_PATH

2. Optimize parameters with COCO validation.

python3 val_optimizer.py --weights ${WEIGHT_PATH | WANDB_PATH} --data-cfg $DATA_CONFIG_PATH --run-json --json-path $JSON_FILE_PATH

The --json-path is optional.

Advanced usage

• If you have baseline network, give --base-map50 and --base-time arguments which are used for objective function.
• To avoid the optimized parameters overfits, use --n-skip option to skip some images.

There are three steps to apply SWA (Stochastic Weight Averaging):

1. Fine-tune pre-trained model
2. Create SWA model
3. Test SWA model

1. Fine-tune pre-trained model

Example

$ python train.py --model yolov5l_kindle.pt \
                  --data res/configs/data/coco.yaml \
                  --cfg res/configs/cfg/finetune.yaml \
                  --wlog --wlog_name yolov5l_swa \
                  --use_swa

2. Create SWA model

Example

$ python create_swa_model.py --model_dir exp/train/2021_1104_runs/weights \
                             --swa_model_name swa_best5.pt \
                             --best_num 5

Usage

$ python create_swa_model.py --help
usage: create_swa_model.py [-h] --model_dir MODEL_DIR
                           [--swa_model_name SWA_MODEL_NAME]
                           [--best_num BEST_NUM]

optional arguments:
  -h, --help            show this help message and exit
  --model_dir MODEL_DIR
                        directory of trained models to apply SWA (default: )
  --swa_model_name SWA_MODEL_NAME
                        file name of SWA model (default: swa.pt)
  --best_num BEST_NUM   the number of trained models to apply SWA (default: 5)

3. Test SWA model

Example

$ python val.py --weights exp/train/2021_1104_runs/weights/swa_best5.pt \
                --model-cfg '' \
                --data-cfg res/configs/data/coco.yaml \
                --conf-t 0.1 --iou-t 0.2