Add train.run() method

[glenn-jocher]

@AyushExel this PR is for passing opt explicitly to train(). It looks like the logging may require significant rework, can you take a look at this please? Thanks!

EDIT: the final result will need to look like detect.py, i.e. with no opt being passed.

# imports here

def train.py(device, hyps, arg1, arg2, etc...)
    # train
    return results


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument(...
    opt = parser.parse_args()
    return opt


def main(opt):
    print(colorstr('train: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
    check_requirements(exclude=('thop'))

    # extra code here
    # extra code here
    # extra code here

    train(device, hyps, **vars(opt))


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Refactoring and optimizations in train.py for enhanced maintainability and readability.

📊 Key Changes

• Variables are now unpacked in a more readable way, with options such as evolve, data, cfg, resume, notest, nosave, workers included in the unpacking sequence.
• Replaced direct references to opt with the unpacked variables for clarity.
• with open(opt.data) changed to with open(data) to use the unpacked data variable.
• Introduced the plots variable to replace not opt.evolve checks.
• Altered TensorBoard and Weights & Biases logging code to utilize the new unpacked variables.
• Refactored model creation code to use the unpacked cfg variable.
• Data loaders and test loaders now use unpacked options for parameters like workers and notest.
• Save conditions at the end of training epochs now use the unpacked nosave and evolve variables.
• Adjusted Weights & Biases logging conditions, changing occurrences of wandb_logger.wandb to loggers['wandb'] for consistency and clarity.
• Simplified the arguments parsing function by allowing for known and custom arguments via parse_opt(known=False).
• Added the run function to enable custom keyword arguments when importing and running the training function from other scripts.

🎯 Purpose & Impact

• Enhanced Readability: Changes improve the code structure, making it easier to follow and maintain.
• Better Logging: Streamlined usage of logging variables for TensorBoard and Weights & Biases.
• Modularity: The run function and improved argument parsing cater to external usage from other scripts, promoting modularity.
• Error Checking: Robust assertion checks and initialization of variables to prevent potential errors during runtime.

🔍 The overall impact is a more maintainable, modular, and error-resistant training script, positively affecting developers and users aiming for better integration and customization.

[AyushExel]

@glenn-jocher thanks. I'll start working based on this branch

[glenn-jocher]

@AyushExel @kalenmike @SkalskiP to take a step back and explain the plan: previously the 4 main functions (train, test, detect, export) were designed for command line ops, i.e. python train.py ..., which meant they were not directly callable in python, i.e. from train import train; train(args...).

So I've been updating all 4 to be python callable in preparation for pip package deployment. 3 are complete: test, detect, export, but train.py has been more complicated. Recently we merged #3687 however, which applies some changes to allow for the argarser defaults to be separately retrieved from python with a small change opt = parser.parse_args() ---> opt = parser.parse_known_args()[0]:

import detect; print(detect.parse_opt())

Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, device='', exist_ok=False, half='client', hide_conf=False, hide_labels=False, imgsz=640, iou_thres=0.45, line_thickness=3, max_det=1000, name='exp', nosave=False, project='runs/detect', save_conf=False, save_crop=False, save_txt=False, source='data/images', update=False, view_img=False, weights='yolov5s.pt')

I'm not sure if this alters our current plan or not. In train.py the opt instance is particular useful as it allows for saving all settings to opt.yaml for use when --resuming a run, and also serves as a single point for W&B/TensorBoard logging(opt). I wonder if we could create an entrypoint for pip that accepts optional arguments and replaces the defaults with them rather than needing to disentangle train() from opt entirely? Perhaps something like this:

def run(**kwargs):
    # Usage: import detect; detect.run(imgsz=320, weights='yolov5m.pt')
    opt = parse_opt()
    for k, v in kwargs.items():
        setattr(opt, k, v)
    main(opt)

To produce this example detect workflow (but using the same strategy for train):

import detect; detect.run(imgsz=320, weights='yolov5m.pt')

detect: weights=yolov5m.pt, source=data/images, imgsz=320, conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False
YOLOv5 🚀 v5.0-218-gc1af67d torch 1.9.0 CPU
Downloading https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5m.pt to yolov5m.pt...
100%|██████████| 41.1M/41.1M [00:01<00:00, 38.3MB/s]

Fusing layers... 
Model Summary: 308 layers, 21356877 parameters, 0 gradients
image 1/2 /Users/glennjocher/PycharmProjects/yolov5/data/images/bus.jpg: 320x256 4 persons, 1 bus, Done. (0.167s)
image 2/2 /Users/glennjocher/PycharmProjects/yolov5/data/images/zidane.jpg: 192x320 2 persons, 2 ties, 1 cell phone, Done. (0.127s)
Results saved to runs/detect/exp6
Done. (0.354s)

EDIT: I implemented my idea in 2b23ea0. Usage: import train; train.run(imgsz=320, weights='yolov5m.pt')

[glenn-jocher]

@AyushExel @kalenmike @SkalskiP I think this should work as an intermediate solution so I'm going to merge this now. Current usage is unchanged, with two new usage examples made possible by 2b23ea0:

git clone

# git clone https://github.com/ultralytics/yolov5 && cd yolov5
import train

train.run(imgsz=320, weights='yolov5m.pt')

pip install

# pip install ultralytics-yolov5
from yolov5 import train

train.run(imgsz=320, weights='yolov5m.pt')

[AyushExel]

@glenn-jocher I'm thinking the easiest approach would be something like this. Let me know if you have other plans.

• train.py changes from def train(opt) to `def train(device, hyps, arg1, arg2, etc...)
• opt = parser.parse_known_args()[0] if known else parser.parse_args()
• convert the arguments to dict -> vars(opt)
• and call training script like this -> train(**opt)

[glenn-jocher]

@AyushExel well I'm cautiously optimistic we've resolved everything for now after this PR. opt is still passed to train() but at least now all the full training commands are python callable, though train is handled a bit differently:

from yolov5 import train, test, detect
from yolov5.models import export

train.run(imgsz=640, data='coco128.yaml')
test.test(imgsz=640, data='coco128.yaml', weights='yolov5s.pt')
detect.detect(imgsz=640)
export.export(imgsz=640, weights='yolov5s.pt')

I guess one option would be to rename test(), detect() and export() functions each to run() (and move export.py to top level) so that all 4 are then imported and called identically? EDIT: implemented in PR #3715

Tomorrow we have a big group meeting and we'll review with @kalenmike and @SkalskiP.