Jupyter notebook written to make process of comparing different YOLO{v5, v6, v7} versions on COCO2017 dataset easier.

yolo-collection-coco2017-train.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# About\nThis notebook enables you to easily train (currently) YOLOv{5,6,7} on MSCOCO2017 dataset.\n\n**Notes:**\n* *Notebook is written with Kaggle in mind, so most paths are absolute and hardcoded for kaggle. You are needed to make necessary changes if running on other platforms.*\n* *Training is set to last only one epoch, you can change that in `python train.py` lines.*\n* *Training, even for one epoch, lasts very long, consider running it in a batch session*\n* *Notebook was tested with GPU P100, if you get `No space left` or `CUDA out of memory`, consider lowering batch size.*\n* *Just starting this notebook without making changes takes about 7hrs.*\n* *After training/testing, all outputs are stored in default output directory.*\n# Directory structure:\n```shell\n/kaggle/\n├─ working/\n│  ├─ yolov5*/\n│  ├─ YOLOv6*/\n│  ├─ yolov7*/\n├─ datasets/ \n│  ├─ coco/\n│  │  ├─ images/\n│  │  ├─ annotations/\n│  │  ├─ labels/\n│  │  ├─ train2017.txt\n│  │  ├─ val2017.txt\n│  │  ├─ test-dev2017.txt\n```\n\\* - optional\n* **","metadata":{}},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nclass YOLO_VERSIONS:\n    YOLOV5 = 'yolov5'\n    YOLOV6 = 'yolov6'\n    YOLOV7 = 'yolov7'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-01-15T21:56:57.040041Z","iopub.execute_input":"2023-01-15T21:56:57.041187Z","iopub.status.idle":"2023-01-15T21:56:57.046916Z","shell.execute_reply.started":"2023-01-15T21:56:57.041142Z","shell.execute_reply":"2023-01-15T21:56:57.045935Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"* **\n# Options\nChoose which **version** of YOLO to use and if you want **evaluation** after testing.\n* **","metadata":{}},{"cell_type":"code","source":"# you may add multiple yolo versions, ex. current_version = [YOLO_VERSIONS.YOLOV5, YOLO_VERSIONS.YOLOV6]\ncurrent_version = [YOLO_VERSIONS.YOLOV5]\neval = True","metadata":{"execution":{"iopub.status.busy":"2023-01-15T21:57:03.582748Z","iopub.execute_input":"2023-01-15T21:57:03.583113Z","iopub.status.idle":"2023-01-15T21:57:03.590961Z","shell.execute_reply.started":"2023-01-15T21:57:03.583084Z","shell.execute_reply":"2023-01-15T21:57:03.588851Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Download coco datasets and takes 30mins+ (~40gb unzipped).\n# NOTE: download directory should not be output, ie /kaggle/working as it's capacity is only ~19gb\n%mkdir -p ../datasets/coco/images\n%cd ../datasets/coco/images\n!wget http://images.cocodataset.org/zips/train2017.zip\nprint(\"Unzipping train dataset...\")\n!unzip -q train2017.zip && rm -rf train2017.zip\n!wget http://images.cocodataset.org/zips/val2017.zip\nprint(\"Unzipping val dataset...\")\n!unzip -q val2017.zip && rm -rf  val2017.zip\n!wget http://images.cocodataset.org/zips/test2017.zip\nprint(\"Unzipping test dataset...\")\n!unzip -q test2017.zip && rm -rf test2017.zip","metadata":{"_kg_hide-output":true,"scrolled":true,"execution":{"iopub.status.busy":"2023-01-15T21:57:09.929464Z","iopub.execute_input":"2023-01-15T21:57:09.929857Z","iopub.status.idle":"2023-01-15T21:57:23.770868Z","shell.execute_reply.started":"2023-01-15T21:57:09.929823Z","shell.execute_reply":"2023-01-15T21:57:23.768642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Downloading COCO labels and annotations\n%cd /kaggle\n%mkdir tmp\n%cd tmp\n!wget https://github.com/meituan/YOLOv6/releases/download/0.1.0/coco2017labels.zip\nprint(\"Unzipping labels and annotations...\")\n!unzip -q coco2017labels.zip && rm -rf coco2017labels.zip\n%cp -r /kaggle/tmp/coco/annotations /kaggle/datasets/coco\n%cp -r /kaggle/tmp/coco/labels /kaggle/datasets/coco\n%cp -r /kaggle/tmp/coco/val2017.txt /kaggle/datasets/coco\n%cp -r /kaggle/tmp/coco/train2017.txt /kaggle/datasets/coco\n%cp -r /kaggle/tmp/coco/test-dev2017.txt /kaggle/datasets/coco\n%pwd\n%ls /kaggle\n%ls /kaggle/working\n%cd /kaggle\n%rm -rf /kaggle/tmp","metadata":{"execution":{"iopub.status.busy":"2023-01-15T21:57:27.094330Z","iopub.execute_input":"2023-01-15T21:57:27.094687Z","iopub.status.idle":"2023-01-15T21:58:12.196919Z","shell.execute_reply.started":"2023-01-15T21:57:27.094658Z","shell.execute_reply":"2023-01-15T21:58:12.195356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Cloning necessary git repos\n\nif YOLO_VERSIONS.YOLOV5 in current_version:\n    %cd /kaggle/working\n    !git clone https://github.com/ultralytics/yolov5  # clone\n    %cd /kaggle/working/yolov5\n    %pip install -r requirements.txt\nif YOLO_VERSIONS.YOLOV6 in current_version:\n    %cd /kaggle/working\n    !git clone https://github.com/meituan/YOLOv6\n    %cd /kaggle/working/YOLOv6\n    %pip install -r requirements.txt\nif YOLO_VERSIONS.YOLOV7 in current_version:\n    %cd /kaggle/working\n    !git clone https://github.com/WongKinYiu/yolov7\n    %cd /kaggle/working/yolov7\n    %pip install -r requirements.txt","metadata":{"_kg_hide-output":true,"scrolled":true,"execution":{"iopub.status.busy":"2023-01-15T22:09:57.532118Z","iopub.execute_input":"2023-01-15T22:09:57.532523Z","iopub.status.idle":"2023-01-15T22:10:10.605057Z","shell.execute_reply.started":"2023-01-15T22:09:57.532492Z","shell.execute_reply":"2023-01-15T22:10:10.603737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"coco_path = '/kaggle/datasets/coco'\n%ls /kaggle/working\n%ls /kaggle\n# change train, val, and test datasets locations in coco.yaml\nif YOLO_VERSIONS.YOLOV5 in current_version:\n    %cp /kaggle/working/yolov5/data/coco.yaml /kaggle/working/yolov5/coco_original.yaml\n    with open('/kaggle/working/yolov5/data/coco.yaml', 'w') as new_file:\n        with open('/kaggle/working/yolov5/coco_original.yaml', 'r') as original_file:\n            for line in original_file:\n                if 'path:' in line:\n                    new_file.write(f'path: {coco_path}\\n')\n                else:\n                    new_file.write(line)\nif YOLO_VERSIONS.YOLOV6 in current_version:\n    %cp /kaggle/working/YOLOv6/data/coco.yaml /kaggle/working/YOLOv6/coco_original.yaml\n    with open('/kaggle/working/YOLOv6/data/coco.yaml', 'w') as new_file:\n        with open('/kaggle/working/YOLOv6/coco_original.yaml', 'r') as original_file:\n            for line in original_file:\n                if 'train:' in line:\n                    new_file.write(f'train: {coco_path}/images/train2017 # 118287 images\\n')\n                elif 'val:' in line:\n                    new_file.write(f'val: {coco_path}/images/val2017   # 5000 images\\n')\n                elif 'test:' in line:\n                    new_file.write(f'test: {coco_path}/images/test2017\\n')\n                elif 'anno_path:' in line:\n                    new_file.write(f'anno_path: {coco_path}/annotations/instances_val2017.json\\n')\n                else:\n                    new_file.write(line)\n                    \nif YOLO_VERSIONS.YOLOV7 in current_version:\n    %cp /kaggle/working/yolov7/data/coco.yaml /kaggle/working/yolov7/coco_original.yaml\n    with open('/kaggle/working/yolov7/data/coco.yaml', 'w') as new_file:\n        with open('/kaggle/working/yolov7/coco_original.yaml', 'r') as original_file:\n            for line in original_file:\n                if 'train:' in line:\n                    new_file.write(f'train: {coco_path}/images/train2017 # 118287 images\\n')\n                elif 'val:' in line:\n                    new_file.write(f'val: {coco_path}/images/val2017  # 5000 images\\n')\n                elif 'test:' in line:\n                    new_file.write(f'test: {coco_path}/test2017\\n')\n                elif 'anno_path:' in line:\n                    new_file.write(f'anno_path: {coco_path}/annotations/instances_val2017.json\\n')\n                else:\n                    new_file.write(line)\n","metadata":{"execution":{"iopub.status.busy":"2023-01-15T22:10:13.092912Z","iopub.execute_input":"2023-01-15T22:10:13.093296Z","iopub.status.idle":"2023-01-15T22:10:16.118101Z","shell.execute_reply.started":"2023-01-15T22:10:13.093262Z","shell.execute_reply":"2023-01-15T22:10:16.116834Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\n# comment this if you want to use wandb (note that in that case wandb-api must be provided [add-ons->secrets])\n%env WANDB_MODE=dryrun\n%pip install pycocotools\nif YOLO_VERSIONS.YOLOV5 in current_version:\n    %cd /kaggle/working/yolov5\n    !python train.py --data coco.yaml --epochs 1 --weights '' --cfg yolov5n.yaml  --batch-size 32\nif YOLO_VERSIONS.YOLOV6 in current_version:\n    %cd /kaggle/working/YOLOv6\n    !python tools/train.py --img 640 --batch 32 --epochs 1 --conf configs/yolov6s.py --data data/coco.yaml\nif YOLO_VERSIONS.YOLOV7 in current_version:\n    %cd /kaggle/working/yolov7\n    # echo 'x' - user choice (replace x with one of the available options)\n    #         wandb: (1) Create a W&B account\n    #         wandb: (2) Use an existing W&B account\n    #         wandb: (3) Don't visualize my results\n    # NOTE: weights flag with value '' means weights are randomly initialized\n    !echo \"3\" | python train.py --workers 8 --device 0 --batch-size 16 --epochs 1 --cfg cfg/training/yolov7.yaml --data data/coco.yaml --img 640 640 --cfg cfg/training/yolov7.yaml --weights '' --name yolov7 --hyp data/hyp.scratch.p5.yaml","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-01-15T22:10:20.544462Z","iopub.execute_input":"2023-01-15T22:10:20.544888Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#eval\n#TODO: requires testing\nif eval:\n    if YOLO_VERSIONS.YOLOV5 in current_version:\n        %cd /kaggle/working/yolov5\n        %cat data/coco.yaml\n        !python val.py --weights /kaggle/working/yolov5/runs/train/exp/weights/best.pt --data coco.yaml --img 640\n    if YOLO_VERSIONS.YOLOV6 in current_version:\n        %cd /kaggle/working/YOLOv6\n        !python tools/eval.py --data data/coco.yaml  --weights /kaggle/working/YOLOv6/runs/train/exp/weights/best_ckpt.pt --task val --device 0\n    if YOLO_VERSIONS.YOLOV7 in current_version:\n        %cd /kaggle/working/yolov7\n        !python test.py --data data/coco.yaml --img 640 --batch 32 --conf 0.001 --iou 0.65 --device 0 --weights /kaggle/working/yolov7/runs/train/yolov7/weights/best.pt --name yolov7_640_val\n","metadata":{"execution":{"iopub.status.busy":"2023-01-15T22:07:41.101845Z","iopub.execute_input":"2023-01-15T22:07:41.102231Z","iopub.status.idle":"2023-01-15T22:07:49.704477Z","shell.execute_reply.started":"2023-01-15T22:07:41.102199Z","shell.execute_reply":"2023-01-15T22:07:49.703285Z"},"trusted":true},"execution_count":null,"outputs":[]}]}