def github(tag='main', deformable=False):
if deformable:
tag = TAGS_DEFORMABLE_DETR[tag]
return hub.github(owner=GITHUB_DEFORMABLE_DETR['owner'],
project=GITHUB_DEFORMABLE_DETR['project'],
tag=tag)
else:
tag = TAGS_DETR[tag]
return hub.github(owner=GITHUB_DETR['owner'],
project=GITHUB_DETR['project'],
tag=tag)
def deploy_yolo5():
parser = argparse.ArgumentParser(
'Deploy a trained YOLO5 checkpoint on S3 by stripping out training states'
)
parser.add_argument('chkpt',
help='Path to a trained checkpoint for deployment')
parser.add_argument('--url', help='s3://bucket/key/to/chkpt.pt')
cfg = parser.parse_args()
from ml import hub
from ml.hub import github
from ml.vision.models.detection.yolo5 import GITHUB
repo = hub.repo(github(**GITHUB), force_reload=False)
chkpt = cfg.chkpt
if chkpt is None:
chkpt = f"{repo}/weights/best.pt"
sys.add_path(repo)
from utils.general import strip_optimizer
before = os.path.getsize(chkpt)
strip_optimizer(chkpt)
after = os.path.getsize(chkpt)
logging.info(
f"Optimized {chkpt} optimized from {before / 1024**2:.2f}MB to {after / 1024**2:.2f}MB"
)
if cfg.url.startswith('s3://'):
logging.info(f"Uploading to {cfg.url}...")
parts = cfg.url[len('s3://'):].split('/')
bucket = parts[0]
key = '/'.join(parts[1:])
hub.upload_s3(chkpt, bucket, key)
else:
ValueError(f"Unsupported URL to upload: {cfg.url}")
def train_yolo5():
"""Train YOLOv5 over a dataset in YOLO5/COCO format.
Usage:
train_yolo5 ../../datasets/Retail81 --with-coco --names object --device 1 --batch-size 32
References:
python train.py --data coco.yaml --cfg yolov5x.yaml --weights yolov5x.pt --device 1 --batch-size 16
"""
parser = argparse.ArgumentParser(
'Train YOLO5 oover a dataset in YOLO5/COCO format')
parser.add_argument('path', help='Path to a dataset in YOLO/COCO format')
parser.add_argument('--arch',
choices=['yolov5l', 'yolov5x'],
default='yolov5x',
help='YOLOv5 model architecture to train')
parser.add_argument('--tag',
choices=['v1.0', 'v2.0', 'v3.0'],
default='v3.0',
help='YOLOv5 repo version tag')
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--device',
default='',
help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--names', nargs='*', default=[], help='class names')
parser.add_argument('--with-coco',
action='store_true',
help='Including COCO80 classes')
args = parser.parse_args()
logging.info(
f"dataset={args.path}, arch={args.arch}, names={args.names}, with_coco={args.with_coco}"
)
path = Path(args.path).resolve()
data = path / 'data'
data.mkdir(exist_ok=True)
data /= f"{path.name.lower()}.yaml"
models = path / "models"
models.mkdir(exist_ok=True)
cfg = f"{args.arch}.yaml"
weights = f"{args.arch}.pt"
names = args.names
if args.with_coco:
from ml.vision.datasets.coco import COCO80_CLASSES
names = COCO80_CLASSES + names
nc = len(names)
assert nc > 0
with open(data, 'w') as f:
yaml = f"""
train: { path / 'train.txt' }
val: { path / 'val.txt' }
nc: {nc}
names: {names}
"""
print(yaml, file=f)
logging.info(yaml)
from ml.vision.models.detection.yolo5 import github
from ml.utils import Config
from ml import hub
repo = hub.repo(github(tag=tag), force_reload=False)
config = Config().load(f"{repo}/models/{cfg}")
config.nc = nc
config.save(models / cfg)
os.chdir(repo)
cmd = f"python train.py --data {data.resolve()} --cfg {models / cfg} --weights {weights} --device {args.device} --batch-size {args.batch_size} --epochs {args.epochs}"
logging.info(f"wd={os.getcwd()}")
logging.info(cmd)
r = os.system(cmd)
if r:
logging.error(f"Failed train YOLOv5 over {path} with res={r}")
else:
# v1.0
# logging.info(f"Best trained checkpoiont at {repo}/weights/best.pt")
# last v1.0 and v2.0+
logging.info(
f"Best trained checkpoiont at {repo}/runs/expX/weights/best.pt")
def github(tag='main'):
tag = TAGS[tag]
return hub.github(owner=GITHUB['owner'],
project=GITHUB['project'],
tag=tag)