def trainer(opt):
# Set DDP variables
opt.total_batch_size = opt.batch_size
opt.world_size = int(
os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
set_logging(opt.global_rank)
if opt.global_rank in [-1, 0]:
check_git_status()
# Resume
if opt.resume: # resume an interrupted run
ckpt = opt.resume if isinstance(
opt.resume,
str) else get_latest_run() # specified or most recent path
log_dir = Path(ckpt).parent.parent # runs/exp0
assert os.path.isfile(
ckpt), 'ERROR: --resume checkpoint does not exist'
with open(log_dir / 'opt.yaml') as f:
opt = argparse.Namespace(**yaml.load(
f, Loader=yaml.FullLoader)) # replace
opt.cfg, opt.weights, opt.resume = '', ckpt, True
logger.info('Resuming training from %s' % ckpt)
else:
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(
opt.cfg), check_file(opt.hyp) # check files
assert len(opt.cfg) or len(
opt.weights), 'either --cfg or --weights must be specified'
opt.img_size.extend(
[opt.img_size[-1]] *
(2 - len(opt.img_size))) # extend to 2 sizes (train, test)
log_dir = increment_dir(Path(opt.logdir) / 'exp',
opt.name) # runs/exp1
device = select_device(opt.device, batch_size=opt.batch_size)
# DDP mode
if opt.local_rank != -1:
assert torch.cuda.device_count() > opt.local_rank
torch.cuda.set_device(opt.local_rank)
device = torch.device('cuda', opt.local_rank)
dist.init_process_group(backend='nccl',
init_method='env://') # distributed backend
assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
opt.batch_size = opt.total_batch_size // opt.world_size
logger.info(opt)
with open(opt.hyp) as f:
hyp = yaml.load(f, Loader=yaml.FullLoader) # load hyps
# Train
if not opt.evolve:
tb_writer = None
if opt.global_rank in [-1, 0]:
logger.info(
'Start Tensorboard with "tensorboard --logdir %s", view at http://localhost:6006/'
% opt.logdir)
tb_writer = SummaryWriter(log_dir=log_dir) # runs/exp0
train(hyp, opt, device, tb_writer)
# Evolve hyperparameters (optional)
else:
# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
meta = {
'lr0':
(1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf':
(1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'giou': (1, 0.02, 0.2), # GIoU loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma':
(0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0,
0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective':
(0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0)
} # image mixup (probability)
assert opt.local_rank == -1, 'DDP mode not implemented for --evolve'
opt.notest, opt.nosave = True, True # only test/save final epoch
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
yaml_file = Path(
'runs/evolve/hyp_evolved.yaml') # save best result here
if opt.bucket:
os.system('gsutil cp gs://%s/evolve.txt .' %
opt.bucket) # download evolve.txt if exists
for _ in range(300): # generations to evolve
if os.path.exists(
'evolve.txt'
): # if evolve.txt exists: select best hyps and mutate
# Select parent(s)
parent = 'single' # parent selection method: 'single' or 'weighted'
x = np.loadtxt('evolve.txt', ndmin=2)
n = min(5, len(x)) # number of previous results to consider
x = x[np.argsort(-fitness(x))][:n] # top n mutations
w = fitness(x) - fitness(x).min() # weights
if parent == 'single' or len(x) == 1:
# x = x[random.randint(0, n - 1)] # random selection
x = x[random.choices(range(n),
weights=w)[0]] # weighted selection
elif parent == 'weighted':
x = (x * w.reshape(
n, 1)).sum(0) / w.sum() # weighted combination
# Mutate
mp, s = 0.8, 0.2 # mutation probability, sigma
npr = np.random
npr.seed(int(time.time()))
g = np.array([x[0] for x in meta.values()]) # gains 0-1
ng = len(meta)
v = np.ones(ng)
while all(
v == 1
): # mutate until a change occurs (prevent duplicates)
v = (g * (npr.random(ng) < mp) * npr.randn(ng) *
npr.random() * s + 1).clip(0.3, 3.0)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# Constrain to limits
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # lower limit
hyp[k] = min(hyp[k], v[2]) # upper limit
hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation
results = train(hyp.copy(), opt, device)
# Write mutation results
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
# Plot results
plot_evolution(yaml_file)
print(
'Hyperparameter evolution complete. Best results saved as: %s\nCommand to train a new model with these '
'hyperparameters: $ python train.py --hyp %s' %
(yaml_file, yaml_file))
npr.seed(int(time.time()))
g = np.array([x[0] for x in meta.values()]) # gains 0-1
ng = len(meta)
v = np.ones(ng)
while all(
v == 1
): # mutate until a change occurs (prevent duplicates)
v = (g * (npr.random(ng) < mp) * npr.randn(ng) *
npr.random() * s + 1).clip(0.3, 3.0)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# Constrain to limits
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # lower limit
hyp[k] = min(hyp[k], v[2]) # upper limit
hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation
results = train(hyp.copy(), opt, device)
# Write mutation results
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
# Plot results
plot_evolution(yaml_file)
print(
'Hyperparameter evolution complete. Best results saved as: %s\nCommand to train a new model with these '
'hyperparameters: $ python train.py --hyp %s' %
(yaml_file, yaml_file))
def main(override_args=None):
parser = argparse.ArgumentParser()
parser.add_argument('--weights',
type=str,
default='',
help='initial weights path')
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
parser.add_argument('--hyp',
type=str,
default='data/hyp.scratch.yaml',
help='hyperparameters path')
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size',
type=int,
default=16,
help='total batch size for all GPUs')
parser.add_argument('--img-size',
nargs='+',
type=int,
default=[640, 640],
help='[train, test] image sizes')
parser.add_argument('--rect',
action='store_true',
help='rectangular training')
parser.add_argument('--resume',
nargs='?',
const=True,
default=False,
help='resume most recent training')
parser.add_argument('--nosave',
action='store_true',
help='only save final checkpoint')
parser.add_argument('--notest',
action='store_true',
help='only test final epoch')
parser.add_argument('--noautoanchor',
action='store_true',
help='disable autoanchor check')
parser.add_argument('--evolve',
action='store_true',
help='evolve hyperparameters')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache-images',
action='store_true',
help='cache images for faster training')
parser.add_argument('--image-weights',
action='store_true',
help='use weighted image selection for training')
parser.add_argument(
'--name',
default='',
help='renames experiment folder exp{N} to exp{N}_{name} if supplied')
parser.add_argument('--device',
default='',
help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale',
action='store_true',
help='vary img-size +/- 50%%')
parser.add_argument('--single-cls',
action='store_true',
help='train as single-class dataset')
parser.add_argument('--adam',
action='store_true',
help='use torch.optim.Adam() optimizer')
parser.add_argument('--sync-bn',
action='store_true',
help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--local_rank',
type=int,
default=-1,
help='DDP parameter, do not modify')
parser.add_argument('--logdir',
type=str,
default='runs/',
help='logging directory')
parser.add_argument('--log-imgs',
type=int,
default=10,
help='number of images for W&B logging, max 100')
parser.add_argument('--workers',
type=int,
default=8,
help='maximum number of dataloader workers')
# changed/added parameters:
parser.add_argument('--data',
type=str,
default='data/coco128.yaml',
help='data.yaml path OR path to the HHI Json Dataset')
# here you can set the path to the "yaml" file that describes the dataset, just as before
# but if you give a path to a JSON file, then it will be read and parsed as a HHI Json Format Dataset metadata file
parser.add_argument(
'--val-size',
type=float,
default=0.05,
help='fraction of the Json Dataset, used for validation')
# if using the HHI Json Format Dataset format, this sets the ratio of the dataset, that will be used for validation (random samples)
parser.add_argument(
'--classes',
type=str,
default='rectangle',
help='limit training to only these class or class list',
nargs="+")
# if using the HHI Json Format Dataset, normally all annotations of the 'rectangle' type will be used for training,
# but you can limit the training to only specific classes, and give them as a space-separated list
# (use ""-quotation marks, if a class name has whitespaces), e.g.:
# python train.py --data ... --classes gec_object bad_gec_object "screw hole" hand "open hand" --val_size 0.05 ...
opt = parser.parse_args()
# ---------------
# override the command-line parameters:
if override_args is not None:
opt.data = override_args.data
opt.cfg = override_args.cfg
opt.multi_scale = override_args.multi_scale
opt.val_size = override_args.val_size
opt.epochs = override_args.epochs
opt.batch_size = override_args.batch_size
opt.classes = override_args.classes
opt.device = override_args.device
# ---------------
# Set DDP variables
opt.total_batch_size = opt.batch_size
opt.world_size = int(
os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
set_logging(opt.global_rank)
if opt.global_rank in [-1, 0]:
check_git_status()
# Resume
if opt.resume: # resume an interrupted run
ckpt = opt.resume if isinstance(
opt.resume,
str) else get_latest_run() # specified or most recent path
log_dir = Path(ckpt).parent.parent # runs/exp0
assert os.path.isfile(
ckpt), 'ERROR: --resume checkpoint does not exist'
with open(log_dir / 'opt.yaml') as f:
opt = argparse.Namespace(**yaml.load(
f, Loader=yaml.FullLoader)) # replace
opt.cfg, opt.weights, opt.resume = '', ckpt, True
logger.info('Resuming training from %s' % ckpt)
else:
# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(
opt.cfg), check_file(opt.hyp) # check files
assert len(opt.cfg) or len(
opt.weights), 'either --cfg or --weights must be specified'
opt.img_size.extend(
[opt.img_size[-1]] *
(2 - len(opt.img_size))) # extend to 2 sizes (train, test)
log_dir = increment_dir(Path(opt.logdir) / 'exp',
opt.name) # runs/exp1
# DDP mode
device = select_device(opt.device, batch_size=opt.batch_size)
if opt.local_rank != -1:
assert torch.cuda.device_count() > opt.local_rank
torch.cuda.set_device(opt.local_rank)
device = torch.device('cuda', opt.local_rank)
dist.init_process_group(backend='nccl',
init_method='env://') # distributed backend
assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
opt.batch_size = opt.total_batch_size // opt.world_size
# Hyperparameters
with open(opt.hyp) as f:
hyp = yaml.load(f, Loader=yaml.FullLoader) # load hyps
if 'box' not in hyp:
warn(
'Compatibility: %s missing "box" which was renamed from "giou" in %s'
% (opt.hyp, 'https://github.com/ultralytics/yolov5/pull/1120'))
hyp['box'] = hyp.pop('giou')
# Train
logger.info(opt)
if not opt.evolve:
tb_writer, wandb = None, None # init loggers
if opt.global_rank in [-1, 0]:
# Tensorboard
logger.info(
f'Start Tensorboard with "tensorboard --logdir {opt.logdir}", view at http://localhost:6006/'
)
tb_writer = SummaryWriter(log_dir=log_dir) # runs/exp0
# W&B
try:
import wandb
assert os.environ.get('WANDB_DISABLED') != 'true'
logger.info(
"Weights & Biases logging enabled, to disable set os.environ['WANDB_DISABLED'] = 'true'"
)
except (ImportError, AssertionError):
opt.log_imgs = 0
logger.info(
"Install Weights & Biases for experiment logging via 'pip install wandb' (recommended)"
)
train(hyp, opt, device, tb_writer, wandb)
# Evolve hyperparameters (optional)
else:
# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
meta = {
'lr0':
(1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf':
(1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'box': (1, 0.02, 0.2), # box loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma':
(0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0,
0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective':
(0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0)
} # image mixup (probability)
assert opt.local_rank == -1, 'DDP mode not implemented for --evolve'
opt.notest, opt.nosave = True, True # only test/save final epoch
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
yaml_file = Path(
opt.logdir
) / 'evolve' / 'hyp_evolved.yaml' # save best result here
if opt.bucket:
os.system('gsutil cp gs://%s/evolve.txt .' %
opt.bucket) # download evolve.txt if exists
for _ in range(300): # generations to evolve
if os.path.exists(
'evolve.txt'
): # if evolve.txt exists: select best hyps and mutate
# Select parent(s)
parent = 'single' # parent selection method: 'single' or 'weighted'
x = np.loadtxt('evolve.txt', ndmin=2)
n = min(5, len(x)) # number of previous results to consider
x = x[np.argsort(-fitness(x))][:n] # top n mutations
w = fitness(x) - fitness(x).min() # weights
if parent == 'single' or len(x) == 1:
# x = x[random.randint(0, n - 1)] # random selection
x = x[random.choices(range(n),
weights=w)[0]] # weighted selection
elif parent == 'weighted':
x = (x * w.reshape(
n, 1)).sum(0) / w.sum() # weighted combination
# Mutate
mp, s = 0.8, 0.2 # mutation probability, sigma
npr = np.random
npr.seed(int(time.time()))
g = np.array([x[0] for x in meta.values()]) # gains 0-1
ng = len(meta)
v = np.ones(ng)
while all(
v == 1
): # mutate until a change occurs (prevent duplicates)
v = (g * (npr.random(ng) < mp) * npr.randn(ng) *
npr.random() * s + 1).clip(0.3, 3.0)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# Constrain to limits
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # lower limit
hyp[k] = min(hyp[k], v[2]) # upper limit
hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation
results = train(hyp.copy(), opt, device)
# Write mutation results
print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
# Plot results
plot_evolution(yaml_file)
print(
f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n'
f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}'
)
