# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

"""

Train a YOLOv5 model on a custom dataset.

Models and datasets download automatically from the latest YOLOv5 release.

Usage - Single-GPU training:

$ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended)

$ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch

Usage - Multi-GPU DDP training:

$ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,3

Models: https://github.com/ultralytics/yolov5/tree/master/models

Datasets: https://github.com/ultralytics/yolov5/tree/master/data

Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

"""

import argparse

import math

import os

import random

import sys

import time

from copy import deepcopy

from datetime import datetime

from pathlib import Path

import numpy as np

import torch

import torch.distributed as dist

import torch.nn as nn

import yaml

from torch.optim import lr_scheduler

from tqdm import tqdm

FILE = Path(__file__).resolve()

ROOT = FILE.parents[0] # YOLOv5 root directory

if str(ROOT) not in sys.path:

sys.path.append(str(ROOT)) # add ROOT to PATH

ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative

import val as validate # for end-of-epoch mAP

from models.experimental import attempt_load

from models.yolo import Model

from utils.autoanchor import check_anchors

from utils.autobatch import check_train_batch_size

from utils.callbacks import Callbacks

from utils.dataloaders import create_dataloader

from utils.downloads import attempt_download, is_url

from utils.general import (LOGGER, check_amp, check_dataset, check_file, check_git_status, check_img_size,

check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path,

init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods,

one_cycle, print_args, print_mutation, strip_optimizer, yaml_save)

from utils.loggers import Loggers

from utils.loggers.wandb.wandb_utils import check_wandb_resume

from utils.loss import ComputeLoss

from utils.metrics import fitness

from utils.plots import plot_evolve

from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,

smart_resume, torch_distributed_zero_first)

LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html

RANK = int(os.getenv('RANK', -1))

WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))

def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary

return results

def parse_opt(known=False):

return parser.parse_known_args()[0] if known else parser.parse_args()

def main(opt, callbacks=Callbacks()):

f'Usage example: $ python train.py --hyp {evolve_yaml}')

def run(**kwargs):

return opt

if __name__ == "__main__":

opt = parse_opt()

main(opt)