def train(hyp, opt, device, tb_writer=None, wandb=None):
logger.info(f'Hyperparameters {hyp}')
log_dir = Path(tb_writer.log_dir) if tb_writer else Path(
opt.logdir) / 'evolve' # logging directory
wdir = log_dir / 'weights' # weights dfirectory
os.makedirs(wdir, exist_ok=True)
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = str(log_dir / 'results.txt')
epochs, batch_size, total_batch_size, weights, rank = \
opt.epochs_init, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Save run settings
with open(log_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(log_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc, names = (1, ['item']) if opt.single_cls else (int(
data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get('anchors'):
ckpt['model'].yaml['anchors'] = round(
hyp['anchors']) # force autoanchor
model = Model(opt.cfg or ckpt['model'].yaml, ch=3,
nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [
] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp[
'lrf']) + hyp['lrf'] # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Logging
if wandb and wandb.run is None:
id = ckpt.get('wandb_id') if 'ckpt' in locals() else None
wandb_run = wandb.init(config=opt,
resume="allow",
project="YOLOv5",
name=os.path.basename(log_dir),
id=id)
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
shutil.copytree(wdir, wdir.parent /
f'weights_backup_epoch{start_epoch - 1}'
) # save previous weights
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
all_test_dataloader = create_dataloader(
test_path,
# '/media/ivan/share/core50_350_1f/test.txt',
imgsz_test,
total_batch_size,
gs,
hyp=hyp,
augment=False,
# cache=opt.cache_images and not opt.notest,
rect=False,
rank=-1,
world_size=opt.world_size,
workers=opt.workers)[0]
root = '/media/ivan/share/core50_350_1f/batches/'
paths = os.listdir(root)
train_paths = []
valid_paths = []
for p in paths:
if 'train' in p:
train_paths.append(root + p)
elif 'val' in p:
valid_paths.append(root + p)
else:
print(p)
extMem = externalMemory(size=200)
print(f'external memory file: {extMem.get_memory_file()}')
train_paths = [
'/media/ivan/share/demoset/train_4.txt',
'/media/ivan/share/demoset/train_2.txt'
]
valid_paths = [
'/media/ivan/share/demoset/valid.txt',
'/media/ivan/share/demoset/valid.txt'
]
# prepare_params(hyp, opt, device, tb_writer=None, wandb=None)
for core_batch in range(2):
print(f'------------CORE50 itertaion №:{core_batch}------------')
train_on_large_batch(core_batch,
train_paths[core_batch],
valid_paths[core_batch],
model,
device,
logger,
imgsz=imgsz,
imgsz_test=imgsz_test,
gs=gs,
opt=opt,
hyp=hyp,
nc=nc,
log_dir=log_dir,
tb_writer=tb_writer,
names=names,
optimizer=optimizer,
extMem=extMem,
scheduler=scheduler,
lf=lf,
best_fitness=best_fitness,
wandb_run=wandb_run)
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
results, maps, times = test.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=model,
single_cls=opt.single_cls,
dataloader=all_test_dataloader,
save_dir=log_dir / 'images' / str(core_batch),
# plots=epoch == 0 or final_epoch, # plot first and last
log_imgs=opt.log_imgs,
verbose=True)
# wandb.log({'per class/AP per class All': maps[0]})
# tb_writer.add_scalar('per class/AP per class All', maps[0])
# Log
tags = [ # train loss
'test/precision', 'test/recall', 'test/mAP_0.5', 'test/mAP_0.5:0.95',
'test/giou_loss', 'test/obj_loss', 'test/cls_loss'
] # params
for x, tag in zip(list(results), tags):
if tb_writer:
tb_writer.add_scalar(tag, x, core_batch) # tensorboard
if wandb:
wandb.log({tag: x}) # W&B
return results
def load_checkpoint(type_,
weights,
device,
cfg=None,
hyp=None,
nc=None,
recipe=None,
resume=None,
rank=-1):
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights[0] if isinstance(weights, list)
or isinstance(weights, tuple) else weights,
map_location=device) # load checkpoint
start_epoch = ckpt['epoch'] + 1 if 'epoch' in ckpt else 0
pickled = isinstance(ckpt['model'], nn.Module)
train_type = type_ == 'train'
ensemble_type = type_ == 'ensemble'
if pickled and ensemble_type:
# load ensemble using pickled
cfg = None
model = attempt_load(weights, map_location=device) # load FP32 model
state_dict = model.state_dict()
else:
# load model from config and weights
cfg = cfg or (ckpt['yaml'] if 'yaml' in ckpt else None) or \
(ckpt['model'].yaml if pickled else None)
model = Model(cfg,
ch=3,
nc=ckpt['nc'] if ('nc' in ckpt and not nc) else nc,
anchors=hyp.get('anchors') if hyp else None).to(device)
model_key = 'ema' if (not train_type and 'ema' in ckpt
and ckpt['ema']) else 'model'
state_dict = ckpt[model_key].float().state_dict(
) if pickled else ckpt[model_key]
# turn gradients for params back on in case they were removed
for p in model.parameters():
p.requires_grad = True
# load sparseml recipe for applying pruning and quantization
recipe = recipe or (ckpt['recipe'] if 'recipe' in ckpt else None)
sparseml_wrapper = SparseMLWrapper(model, recipe)
exclude_anchors = train_type and (cfg or hyp.get('anchors')) and not resume
loaded = False
if not train_type:
# apply the recipe to create the final state of the model when not training
sparseml_wrapper.apply()
else:
# intialize the recipe for training and restore the weights before if no quantized weights
quantized_state_dict = any(
[name.endswith('.zero_point') for name in state_dict.keys()])
if not quantized_state_dict:
state_dict = load_state_dict(model,
state_dict,
train=True,
exclude_anchors=exclude_anchors)
loaded = True
sparseml_wrapper.initialize(start_epoch)
if not loaded:
state_dict = load_state_dict(model,
state_dict,
train=train_type,
exclude_anchors=exclude_anchors)
model.float()
report = 'Transferred %g/%g items from %s' % (
len(state_dict), len(model.state_dict()), weights)
return model, {
'ckpt': ckpt,
'state_dict': state_dict,
'start_epoch': start_epoch,
'sparseml_wrapper': sparseml_wrapper,
'report': report,
}
def train(hyp, opt, device, tb_writer=None):
print(f'Hyperparameters {hyp}')
log_dir = tb_writer.log_dir if tb_writer else 'runs/evolve' # run directory
wdir = str(Path(log_dir) / 'weights') + os.sep # weights directory
os.makedirs(wdir, exist_ok=True)
last = wdir + 'last.pt'
best = wdir + 'best.pt'
results_file = log_dir + os.sep + 'results.txt'
epochs, batch_size, total_batch_size, weights, rank = \
opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.local_rank
# Save run settings
with open(Path(log_dir) / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
# train_path = data_dict['train']
train_path = f'{opt.trainset_path}/images/train'
# test_path = data_dict['val']
test_path = f'{opt.trainset_path}/images/test'
nc, names = (1, ['item']) if opt.single_cls else (int(
data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Remove previous results
if rank in [-1, 0]:
for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):
os.remove(f)
# Create model
if opt.not_use_SE:
model = Model(opt.cfg, nc=nc).to(device)
else:
model = ModelSE(opt.cfg, nc=nc).to(device)
print(model)
exit()
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# Optimizer
nbs = 64 # nominal batch size
# default DDP implementation is slow for accumulation according to: https://pytorch.org/docs/stable/notes/ddp.html
# all-reduce operation is carried out during loss.backward().
# Thus, there would be redundant all-reduce communications in a accumulation procedure,
# which means, the result is still right but the training speed gets slower.
# in https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/run_pretraining.py
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
if v.requires_grad:
if '.bias' in k:
pg2.append(v) # biases
elif '.weight' in k and '.bn' not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
print('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = lambda x: ((
(1 + math.cos(x * math.pi / epochs)) / 2)**1.0) * 0.8 + 0.2 # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Load Model
with torch_distributed_zero_first(rank):
attempt_download(weights)
start_epoch, best_fitness = 0, 0.0
if weights.endswith('.pt'): # pytorch format
ckpt = torch.load(weights, map_location=device) # load checkpoint
# load model
try:
exclude = ['anchor'] # exclude keys
ckpt['model'] = {
k: v
for k, v in ckpt['model'].float().state_dict().items()
if k in model.state_dict() and not any(x in k for x in exclude)
and model.state_dict()[k].shape == v.shape
}
model.load_state_dict(ckpt['model'], strict=False)
print('Transferred %g/%g items from %s' %
(len(ckpt['model']), len(model.state_dict()), weights))
except KeyError as e:
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \
"Please delete or update %s and try again, or use --weights '' to train from scratch." \
% (weights, opt.cfg, weights, weights)
raise KeyError(s) from e
# load optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# load results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# epochs
start_epoch = ckpt['epoch'] + 1
if epochs < start_epoch:
print(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt
# DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
# model = torch.nn.DataParallel(model)
model = torch.nn.DataParallel(model, device_ids=[0, 1])
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
print('Using SyncBatchNorm()')
# Exponential moving average
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if cuda and rank != -1:
model = DDP(model, device_ids=[rank], output_device=rank)
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
local_rank=rank,
world_size=opt.world_size)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# Testloaderc
if rank in [-1, 0]:
# local_rank is set to -1. Because only the first process is expected to do evaluation.
testloader = create_dataloader(test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images,
rect=True,
local_rank=-1,
world_size=opt.world_size)[0]
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device) # attach class weights
model.names = names
# Class frequency
if rank in [-1, 0]:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1.
# model._initialize_biases(cf.to(device))
plot_labels(labels, save_dir=log_dir)
if tb_writer:
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
tb_writer.add_histogram('classes', c, 0)
# Check anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# Start training
t0 = time.time()
nw = max(3 * nb,
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (
0, 0, 0, 0, 0, 0, 0
) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
if rank in [0, -1]:
print('Image sizes %g train, %g test' % (imgsz, imgsz_test))
print('Using %g dataloader workers' % dataloader.num_workers)
print('Starting training for %g epochs...' % epochs)
# torch.autograd.set_detect_anomaly(True)
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if dataset.image_weights:
# Generate indices
if rank in [-1, 0]:
w = model.class_weights.cpu().numpy() * (
1 - maps)**2 # class weights
image_weights = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=w)
dataset.indices = random.choices(
range(dataset.n), weights=image_weights,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = torch.zeros([dataset.n], dtype=torch.int)
if rank == 0:
indices[:] = torch.from_tensor(dataset.indices,
dtype=torch.int)
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
if rank in [-1, 0]:
print(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj',
'cls', 'total', 'targets', 'img_size'))
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(
ni, xi,
[0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi,
[0.9, hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Autocast
with amp.autocast(enabled=cuda):
# Forward
pred = model(imgs)
# print([x.shape for x in pred]) # [1, 3, 76, 76, 25] [1, 3, 38, 38, 25] [1, 3, 19, 19, 25])
# Loss
loss, loss_items = compute_loss(pred, targets.to(device),
model) # scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
# if not torch.isfinite(loss):
# print('WARNING: non-finite loss, ending training ', loss_items)
# return results
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema is not None:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if ni < 3:
f = str(Path(log_dir) /
('train_batch%g.jpg' % ni)) # filename
result = plot_images(images=imgs,
targets=targets,
paths=paths,
fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f,
result,
dataformats='HWC',
global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema is not None:
ema.update_attr(
model,
include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
save_json=final_epoch
and opt.data.endswith(os.sep + 'coco.yaml'),
model=ema.ema.module
if hasattr(ema.ema, 'module') else ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir)
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results +
'\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' %
(results_file, opt.bucket, opt.name))
# Tensorboard
if tb_writer:
tags = [
'train/giou_loss', 'train/obj_loss', 'train/cls_loss',
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95', 'val/giou_loss', 'val/obj_loss',
'val/cls_loss'
]
for x, tag in zip(list(mloss[:-1]) + list(results), tags):
tb_writer.add_scalar(tag, x, epoch)
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
f.read(),
'model':
ema.ema.module if hasattr(ema, 'module') else ema.ema,
'optimizer':
None if final_epoch else optimizer.state_dict()
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
n = ('_'
if len(opt.name) and not opt.name.isnumeric() else '') + opt.name
fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n
for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'],
[flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
ispt = f2.endswith('.pt') # is *.pt
strip_optimizer(f2) if ispt else None # strip optimizer
os.system('gsutil cp %s gs://%s/weights' % (
f2, opt.bucket)) if opt.bucket and ispt else None # upload
# Finish
if not opt.evolve:
plot_results(save_dir=log_dir) # save as results.png
print('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
return results
def train(
hyp, # path/to/hyp.yaml or hyp dictionary
opt,
device,
):
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, notest, nosave, workers, = \
opt.save_dir, opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
opt.resume, opt.notest, opt.nosave, opt.workers
# Directories
save_dir = Path(save_dir)
wdir = save_dir / 'weights'
wdir.mkdir(parents=True, exist_ok=True) # make dir
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = save_dir / 'results.txt'
# Hyperparameters
if isinstance(hyp, str):
with open(hyp) as f:
hyp = yaml.safe_load(f) # load hyps dict
logger.info(
colorstr('hyperparameters: ') + ', '.join(f'{k}={v}'
for k, v in hyp.items()))
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.safe_dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.safe_dump(vars(opt), f, sort_keys=False)
# Configure
plots = not evolve # create plots
cuda = device.type != 'cpu'
init_seeds(2 + RANK)
with open(data) as f:
data_dict = yaml.safe_load(f) # data dict
# Loggers
loggers = {'wandb': None, 'tb': None} # loggers dict
if RANK in [-1, 0]:
# TensorBoard
if not evolve:
prefix = colorstr('tensorboard: ')
logger.info(
f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/"
)
loggers['tb'] = SummaryWriter(str(save_dir))
# W&B
opt.hyp = hyp # add hyperparameters
run_id = torch.load(weights).get('wandb_id') if weights.endswith(
'.pt') and os.path.isfile(weights) else None
run_id = run_id if opt.resume else None # start fresh run if transfer learning
wandb_logger = WandbLogger(opt, save_dir.stem, run_id, data_dict)
loggers['wandb'] = wandb_logger.wandb
if loggers['wandb']:
data_dict = wandb_logger.data_dict
weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp # may update weights, epochs if resuming
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
names = ['item'] if single_cls and len(
data_dict['names']) != 1 else data_dict['names'] # class names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, data) # check
is_coco = data.endswith('coco.yaml') and nc == 80 # COCO dataset
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(RANK):
weights = attempt_download(
weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
model = Model(cfg or ckpt['model'].yaml,
ch=3,
nc=nc,
anchors=hyp.get('anchors')).to(device) # create
exclude = [
'anchor'
] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(cfg, ch=3, nc=nc,
anchors=hyp.get('anchors')).to(device) # create
with torch_distributed_zero_first(RANK):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp[
'lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA
ema = ModelEMA(model) if RANK in [-1, 0] else None
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Results
if ckpt.get('training_results') is not None:
results_file.write_text(
ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
if opt.stop:
epochs = start_epoch + 1
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = max(int(model.stride.max()), 32) # grid size (max stride)
nl = model.model[
-1].nl # number of detection layers (used for scaling hyp['obj'])
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
logging.warning(
'DP not recommended, instead use torch.distributed.run for best DDP Multi-GPU results.\n'
'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.'
)
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '))
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, data, nc - 1)
# Process 0
if RANK in [-1, 0]:
testloader = create_dataloader(test_path,
imgsz_test,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=opt.cache_images and not notest,
rect=True,
rank=-1,
workers=workers,
pad=0.5,
prefix=colorstr('val: '))[0]
if not resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
if plots:
plot_labels(labels, names, save_dir, loggers)
if loggers['tb']:
loggers['tb'].add_histogram('classes', c, 0) # TensorBoard
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
model.half().float() # pre-reduce anchor precision
# DDP mode
if cuda and RANK != -1:
model = DDP(
model,
device_ids=[LOCAL_RANK],
output_device=LOCAL_RANK,
# nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
find_unused_parameters=any(
isinstance(layer, nn.MultiheadAttention)
for layer in model.modules()))
# Model parameters
hyp['box'] *= 3. / nl # scale to layers
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640)**2 * 3. / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(
dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Start training
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * nb),
1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
compute_loss = ComputeLoss(model) # init loss class
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
f'Using {dataloader.num_workers} dataloader workers\n'
f'Logging results to {save_dir}\n'
f'Starting training for {epochs} epochs...')
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if RANK in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (
1 - maps)**2 / nc # class weights
iw = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=cw) # image weights
dataset.indices = random.choices(
range(dataset.n), weights=iw,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if RANK != -1:
indices = (torch.tensor(dataset.indices)
if RANK == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if RANK != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if RANK != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls',
'total', 'labels', 'img_size'))
if RANK in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if RANK in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % (f'{epoch}/{epochs - 1}', mem, *mloss,
targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if plots and ni < 3:
f = save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images,
args=(imgs, targets, paths, f),
daemon=True).start()
if loggers['tb'] and ni == 0: # TensorBoard
with warnings.catch_warnings():
warnings.simplefilter(
'ignore') # suppress jit trace warning
loggers['tb'].add_graph(
torch.jit.trace(de_parallel(model),
imgs[0:1],
strict=False), [])
elif plots and ni == 10 and loggers['wandb']:
wandb_logger.log({
'Mosaics': [
loggers['wandb'].Image(str(x), caption=x.name)
for x in save_dir.glob('train*.jpg') if x.exists()
]
})
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
# DDP process 0 or single-GPU
if RANK in [-1, 0]:
# mAP
ema.update_attr(model,
include=[
'yaml', 'nc', 'hyp', 'gr', 'names', 'stride',
'class_weights'
])
final_epoch = epoch + 1 == epochs
if not notest or final_epoch: # Calculate mAP
wandb_logger.current_epoch = epoch + 1
results, maps, _ = test.run(data_dict,
batch_size=batch_size //
WORLD_SIZE * 2,
imgsz=imgsz_test,
model=ema.ema,
single_cls=single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=is_coco and final_epoch,
verbose=nc < 50 and final_epoch,
plots=plots and final_epoch,
wandb_logger=wandb_logger,
compute_loss=compute_loss)
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results +
'\n') # append metrics, val_loss
# Log
tags = [
'train/box_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if loggers['tb']:
loggers['tb'].add_scalar(tag, x, epoch) # TensorBoard
if loggers['wandb']:
wandb_logger.log({tag: x}) # W&B
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
wandb_logger.end_epoch(best_result=best_fitness == fi)
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
results_file.read_text(),
'model':
deepcopy(de_parallel(model)).half(),
'ema':
deepcopy(ema.ema).half(),
'updates':
ema.updates,
'optimizer':
optimizer.state_dict(),
'wandb_id':
wandb_logger.wandb_run.id if loggers['wandb'] else None
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if loggers['wandb']:
if ((epoch + 1) % opt.save_period == 0
and not final_epoch) and opt.save_period != -1:
wandb_logger.log_model(last.parent,
opt,
epoch,
fi,
best_model=best_fitness == fi)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training -----------------------------------------------------------------------------------------------------
if RANK in [-1, 0]:
logger.info(
f'{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.\n'
)
if plots:
plot_results(save_dir=save_dir) # save as results.png
if loggers['wandb']:
files = [
'results.png', 'confusion_matrix.png',
*[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]
]
wandb_logger.log({
"Results": [
loggers['wandb'].Image(str(save_dir / f), caption=f)
for f in files if (save_dir / f).exists()
]
})
if not evolve:
if is_coco: # COCO dataset
for m in [last, best
] if best.exists() else [last]: # speed, mAP tests
results, _, _ = test.run(
data,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz_test,
conf_thres=0.001,
iou_thres=0.7,
model=attempt_load(m, device).half(),
single_cls=single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=True,
plots=False)
# Strip optimizers
for f in last, best:
if f.exists():
strip_optimizer(f) # strip optimizers
if loggers['wandb']: # Log the stripped model
loggers['wandb'].log_artifact(
str(best if best.exists() else last),
type='model',
name='run_' + wandb_logger.wandb_run.id + '_model',
aliases=['latest', 'best', 'stripped'])
wandb_logger.finish_run()
torch.cuda.empty_cache()
return results
nargs='+',
type=int,
default=[640, 640],
help='image size')
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
print(opt)
set_logging()
# Input
img = torch.zeros((opt.batch_size, 3,
*opt.img_size)) # image size(1,3,320,192) iDetection
# Load PyTorch model
attempt_download(opt.weights)
model = torch.load(opt.weights,
map_location=torch.device('cpu'))['model'].float()
model.eval()
model.model[-1].export = True # set Detect() layer export=True
y = model(img) # dry run
# TorchScript export
try:
print('\nStarting TorchScript export with torch %s...' %
torch.__version__)
f = opt.weights.replace('.pt', '.torchscript.pt') # filename
ts = torch.jit.trace(model, img)
ts.save(f)
print('TorchScript export success, saved as %s' % f)
except Exception as e:
def train(hyp, opt, device, tb_writer=None):
logger.info(f'Hyperparameters {hyp}')
print("tb_writer.log_dir: ", tb_writer.log_dir) #resumed的时候可以从这里恢复
log_dir = Path(tb_writer.log_dir) if tb_writer else Path(
opt.logdir) / 'evolve' # logging directory
wdir = log_dir / 'weights' # weights directory
os.makedirs(wdir, exist_ok=True)
last = wdir / 'last.pt'
best = wdir / 'best.pt'
print("log_dir: ", log_dir)
print("wdir: ", wdir)
print("last: ", last)
print("best: ", best)
results_file = str(log_dir / 'results.txt')
epochs, batch_size, total_batch_size, weights, rank = \
opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
print("epochs: ", epochs)
print("batch_size: ", batch_size)
print("total_batch_size: ", total_batch_size)
print("weights: ", weights)
print("rank: ", rank)
# Save run settings
with open(log_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(log_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
cuda = device.type != 'cpu'
print("cuda: ", cuda)
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
print("data_dict: ", data_dict)
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc, names = (1, ['item']) if opt.single_cls else (int(
data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
print("pretrained: ", pretrained)
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
"""ckpt中，['epoch'], ['best_fitness'], ['training_results'], ['model'], ['optimizer']"""
# print("ckpt: ", ckpt)
print(
"ckpt: ['epoch'], ['best_fitness'], ['training_results'], ['optimizer']: ",
ckpt['epoch'], ckpt['best_fitness'], ckpt['training_results'],
ckpt['optimizer'])
# print("ckpt['model']: ", ckpt['model'])
# print("ckpt['model'].model: ", ckpt['model'].model)
# print("ckpt['model'].state_dict(): ", ckpt['model'].state_dict())
print("ckpt['model'].save: ", ckpt['model'].save)
print("ckpt['model'].yaml: ", ckpt['model'].yaml)
print("hyp.get('anchors'): ", hyp.get('anchors'))
if hyp.get('anchors'):
ckpt['model'].yaml['anchors'] = round(
hyp['anchors']) # force autoanchor
print("opt.cfg: ", opt.cfg)
# create, 都为真则取or前面的, 即以opt.cfg中的内容(eg:yolov5s.yaml, yolov5x.yaml...)为主，其次是ckpt['model'].yaml
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device)
exclude = (['anchor'] if opt.cfg or hyp.get('anchors') else []
) # exclude keys, 以opt.cfg中的anchor为主，其次是hyp中的
print("exclude: ", exclude)
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
# print("state_dict: ", state_dict)
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = [
'',
] # parameter names to freeze (full or partial)
if any(freeze):
for k, v in model.named_parameters():
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
"""当模型梯度累积了(nbs/total_batch_size)次之后，再更新一次模型参数，变相的扩大了batch_size"""
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
v.requires_grad = True
if '.bias' in k:
pg2.append(v) # biases
elif '.weight' in k and '.bn' not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
#optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# 设置学习率衰减,这里为余弦退火方式进行衰减
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp[
'lrf']) + hyp['lrf'] # cosine
scheduler = lr_scheduler.LambdaLR(optimizer,
lr_lambda=lf) #每个epochs具有一个不同的学习率
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
shutil.copytree(wdir, wdir.parent /
f'weights_backup_epoch{start_epoch - 1}'
) # save previous weights
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride) # 获取模型总步长和模型输入图片分辨率
imgsz, imgsz_test = [
check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples #检查输入图片分辨率是gs=32的整数倍
print("imgsz: ", imgsz)
print("imgsz_test: ", imgsz_test)
# DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
print(
"DP mode..............................................................................."
)
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# Exponential moving average
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if cuda and rank != -1:
model = DDP(model,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers)
print("dataloader: ", dataloader)
print("dataset: ", dataset)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
nb = len(dataloader) # number of batches
print("nb: ", nb)
# Process 0
if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates
print("ema.updates: ", ema.updates)
testloader = create_dataloader(test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images
and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers)[0] # testloader
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
plot_labels(labels, save_dir=log_dir)
if tb_writer:
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
tb_writer.add_histogram('classes', c, 0)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device) # attach class weights
model.names = names
#print("model: ", model)
# Start training
t0 = time.time()
n_warmup = max(
round(hyp['warmup_epochs'] * nb),
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
# n_warmup = min(n_warmup, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
print("n_warmup: ", n_warmup)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda) #混合精度梯度放大模块
maps = np.zeros(nc) # mAP per class
results = (
0, 0, 0, 0, 0, 0, 0
) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
logger.info(
'Image sizes %g train, %g test\nUsing %g dataloader workers\nLogging results to %s\n'
'Starting training for %g epochs...' %
(imgsz, imgsz_test, dataloader.num_workers, log_dir, epochs))
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if rank in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (
1 - maps)**2 # class weights
iw = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=cw) # image weights
dataset.indices = random.choices(
range(dataset.n), weights=iw,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = (torch.tensor(dataset.indices)
if rank == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls',
'total', 'targets', 'img_size'))
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start batch)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= n_warmup:
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(
1,
np.interp(ni, [0, n_warmup],
[1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, [0, n_warmup], [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, [0, n_warmup],
[hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
with amp.autocast(enabled=cuda): #自动混合精度
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device),
model) # loss scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
# Backward
scaler.scale(loss).backward()
# Optimize
"""每accumulate个batch时更新一次, 在n_warmup之内时，accumulate从1逐渐增大到4"""
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if ni < 3:
f = str(log_dir / ('train_batch%g.jpg' % ni)) # filename
result = plot_images(images=imgs,
targets=targets,
paths=paths,
fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f,
result,
dataformats='HWC',
global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema:
ema.update_attr(
model,
include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
if final_epoch: # replot predictions
[
os.remove(x) for x in glob.glob(
str(log_dir / 'test_batch*_pred.jpg'))
if os.path.exists(x)
]
results, maps, times = evaluate.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir)
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results +
'\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' %
(results_file, opt.bucket, opt.name))
# Tensorboard
if tb_writer:
tags = [
'train/giou_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/giou_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
tb_writer.add_scalar(tag, x, epoch)
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
f.read(),
'model':
ema.ema,
'optimizer':
None if final_epoch else optimizer.state_dict()
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
n = opt.name if opt.name.isnumeric() else ''
fresults, flast, fbest = log_dir / f'results{n}.txt', wdir / f'last{n}.pt', wdir / f'best{n}.pt'
for f1, f2 in zip([wdir / 'last.pt', wdir / 'best.pt', results_file],
[flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
if str(f2).endswith('.pt'): # is *.pt
strip_optimizer(f2) # strip optimizer
os.system(
'gsutil cp %s gs://%s/weights' %
(f2, opt.bucket)) if opt.bucket else None # upload
# Finish
if not opt.evolve:
plot_results(save_dir=log_dir) # save as results.png
logger.info('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
return results
def _create(name,
pretrained=True,
channels=3,
classes=80,
autoshape=True,
verbose=True):
"""Creates a specified YOLOv5 model
Arguments:
name (str): name of model, i.e. 'yolov5s'
pretrained (bool): load pretrained weights into the model
channels (int): number of input channels
classes (int): number of model classes
autoshape (bool): apply YOLOv5 .autoshape() wrapper to model
verbose (bool): print all information to screen
Returns:
YOLOv5 pytorch model
"""
from pathlib import Path
from models.yolo import Model, attempt_load
from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device
check_requirements(Path(__file__).parent / 'requirements.txt',
exclude=('tensorboard', 'pycocotools', 'thop'))
set_logging(verbose=verbose)
fname = Path(name).with_suffix('.pt') # checkpoint filename
try:
if pretrained and channels == 3 and classes == 80:
model = attempt_load(
fname,
map_location=torch.device('cpu')) # download/load FP32 model
else:
cfg = list((Path(__file__).parent /
'models').rglob(f'{name}.yaml'))[0] # model.yaml path
model = Model(cfg, channels, classes) # create model
if pretrained:
attempt_download(fname) # download if not found locally
ckpt = torch.load(fname,
map_location=torch.device('cpu')) # load
msd = model.state_dict() # model state_dict
csd = ckpt['model'].float().state_dict(
) # checkpoint state_dict as FP32
csd = {
k: v
for k, v in csd.items() if msd[k].shape == v.shape
} # filter
model.load_state_dict(csd, strict=False) # load
if len(ckpt['model'].names) == classes:
model.names = ckpt[
'model'].names # set class names attribute
if autoshape:
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
device = select_device('0' if torch.cuda.is_available() else
'cpu') # default to GPU if available
return model.to(device)
except Exception as e:
help_url = 'https://github.com/ultralytics/yolov5/issues/36'
s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url
raise Exception(s) from e
def detect(save_img=False):
## Writing output video
fps = 30
codec = cv2.VideoWriter_fourcc(*'XVID')
output_path = './name_video.avi'
out_vid_write = cv2.VideoWriter(output_path, codec, fps, (width, height))
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
fps_disp = 0.0
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
## Init
first_detection = True
lmk_id = 0
old_objects = {}
old_num = 0
codebook_match = {}
txt = " "
correct_hist = {}
rbt_glb_pose = robot_global_pose() ## robot global pose
lmk_global = landmarks_global_pose() ## landmarks global pose
lmk_obsv_poses = {}
spatial_filter_text = ' '
frames_wo_det = 0
## Main loop
for path, img, im0s, d_map, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
dmap = d_map
if dmap is not None:
d4d = np.uint8(dmap.astype(float) * 255 / 2**12 -
1) # Correct the range. Depth images are 12bits
d4d = 255 - cv2.cvtColor(d4d, cv2.COLOR_GRAY2RGB)
timestamp = rospy.Time.now()
rbt_glb_pose.trans_mat()
lmk_global.lmk_gPoses_list()
if img.ndimension() == 3:
img = img.unsqueeze(0)
## Online operation data
online_data = {
'timestamp': timestamp,
'depth_map': dmap,
'robot_global_pose': rbt_glb_pose.trans
}
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
im_rgb = cv2.cvtColor(im0, cv2.COLOR_BGR2RGB)
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(save_path[:save_path.rfind('.')] + '.txt',
'a') as file:
file.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
#plot_one_box(xyxy, im_rgb, label=label, color=colors[int(cls)], line_thickness=3)
plot_one_box(xyxy,
im_rgb,
label=label,
color=(0, 255, 0),
line_thickness=3)
## Check if any detections happen
detections = pred[0]
objects = detections
if det is not None:
## number of detections:
num_d = len(detections)
for i in range(num_d):
det_info = detections[i].data.cpu().tolist()
bbox = [int(j)
for j in det_info[0:4]] ## bbox = [x1 y1 x2 y2]
## Check if the object is within range
obj_cent = [
(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
] ## [x y] center pixel of the area the object is located
## Optimal operation range: 0.8-3.5m
im_rgb = cv2.circle(im_rgb,
(int(obj_cent[0]), int(obj_cent[1])),
5, (0, 255, 0), -1)
d4d = cv2.circle(d4d, (int(obj_cent[0]), int(obj_cent[1])),
5, (0, 0, 255), -1)
if dmap[int(obj_cent[1]),
int(obj_cent[0])] < 800 or dmap[
int(obj_cent[1]),
int(obj_cent[0])] > 3500:
objects = torch.cat(
[detections[0:i], detections[i + 1:]])
if not objects.cpu().tolist(): objects = None
## Track and sample landmarks
if objects is not None:
if first_detection:
frames_wo_det = 0
lmk_id += 1
old_num = len(objects)
im_rgb, old_objects, tracked_histograms, lmkObsv = new_landmarks(
objects, lmk_id, im0, im_rgb, online_data, names)
first_detection = False
codebook_match = {}
correct_hist = {}
else:
## Track objects and serch for best match in keyframes
tracker = track_detections(old_num, im0, im_rgb,
old_objects, objects, lmk_id,
tracked_histograms,
codebook_match, correct_hist,
online_data, lmkObsv, names,
frames_wo_det)
lmk_id = tracker.id
old_objects = tracker.old_objects
old_num = tracker.old_num
im_rgb = tracker.disp
tracked_histograms = tracker.tracked_histograms
codebook_match = tracker.codebook_match
correct_hist = tracker.keyframes_hist
lmkObsv = tracker.lmkObsv
frames_wo_det = 0
if tracker.publish_flag:
## Filter spatially and publish correct landmarks to cartographer
lmk_gp = lmk_global.landmarks_global_poses ## Global poses of landmarks
filtered_obsv = spatial_filter(lmkObsv, correct_hist,
tracked_histograms,
lmk_gp, lmk_obsv_poses,
spatial_filter_text,
lmk_pub, lmk_id)
correct_hist = filtered_obsv.keyframes_copy
tracked_histograms = filtered_obsv.tracked_hists
lmkObsv = filtered_obsv.lmkObsv
spatial_filter_text = filtered_obsv.spatial_filter_text
lmk_id = filtered_obsv.lmk_id
if hasattr(tracker,
'sampler_txt') and tracker.sampler_txt != " ":
txt = tracker.sampler_txt
else: ## no detections
frames_wo_det += 1
img2 = cv2.putText(im_rgb, txt, (0, 60),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 1,
cv2.LINE_AA)
img2 = cv2.putText(im_rgb, spatial_filter_text, (0, 450),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (255, 0, 0), 1,
cv2.LINE_AA)
fps_disp = (fps_disp + (1. /
(torch_utils.time_synchronized() - t1))) / 2
img2 = cv2.putText(img2, "FPS: {:.2f}".format(fps_disp), (0, 30),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
# Stream results
if view_img or save_img:
out_vid_write.write(img2)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, display)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps,
(w, h))
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def load_checkpoint(type_,
weights,
device,
cfg=None,
hyp=None,
nc=None,
recipe=None,
resume=None,
rank=-1):
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
start_epoch = ckpt['epoch'] + 1 if 'epoch' in ckpt else 0
pickled = isinstance(ckpt['model'], nn.Module)
if pickled and type_ == 'ensemble':
# load ensemble using pickled
cfg = None
model = attempt_load(weights, map_location=device) # load FP32 model
state_dict = model.state_dict()
else:
# load model from config and weights
cfg = cfg or (ckpt['yaml'] if 'yaml' in ckpt else None) or \
(ckpt['model'].yaml if pickled else None)
model = Model(cfg,
ch=3,
nc=ckpt['nc'] if ('nc' in ckpt and not nc) else nc,
anchors=hyp.get('anchors') if hyp else None).to(device)
model_key = 'ema' if (type_ in ['ema', 'ensemble'] and 'ema' in ckpt
and ckpt['ema']) else 'model'
state_dict = ckpt[model_key].float().state_dict(
) if pickled else ckpt[model_key]
# turn gradients for params back on in case they were removed
for p in model.parameters():
p.requires_grad = True
# load sparseml recipe for applying pruning and quantization
recipe = recipe or (ckpt['recipe'] if 'recipe' in ckpt else None)
sparseml_wrapper = SparseMLWrapper(model, recipe)
if type_ in ['ema', 'ensemble']:
# apply the recipe to create the final state of the model when not training
sparseml_wrapper.apply()
else:
# intialize the recipe for training
sparseml_wrapper.initialize(start_epoch)
if type_ == 'train':
# load any missing weights from the model
exclude = [
'anchor'
] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=type_ != 'train') # load
model.float()
report = 'Transferred %g/%g items from %s' % (
len(state_dict), len(model.state_dict()), weights)
return model, {
'ckpt': ckpt,
'state_dict': state_dict,
'start_epoch': start_epoch,
'sparseml_wrapper': sparseml_wrapper,
'report': report,
}
def train(hyp, opt, device, tb_writer=None):
# 控制台打印日志
logger.info(f'Hyperparameters {hyp}')
log_dir = Path(tb_writer.log_dir) if tb_writer else Path(
opt.logdir) / 'evolve' # logging directory
wdir = log_dir / 'weights' # weights directory
os.makedirs(wdir, exist_ok=True)
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = str(log_dir / 'results.txt')
# weights：权重文件（预训练的）；rank：全局进程；
epochs, batch_size, total_batch_size, weights, rank = \
opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Save run settings
with open(log_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(log_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
cuda = device.type != 'cpu'
# 初始化随机种子（numpy，random，torch的）
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
'''
执行逻辑：如果是0号进程进来会直接从torch_distributed_zero_first返回,然后执行check_dataset，
执行完check_dataset后会进入torch_distributed_zero_first函数从yield之下开始执行barrier函数暂停等到所有线程都到这个函数再继续执行，
如果不是0号线程进入那么就会执行barrier函数等待，等到所有进程都进入此函数的时候解除barrier继续执行
解除barrier的方法就是等所有进程都执行barrier函数的时候就会解除
'''
# 核实数据
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc, names = (1, ['item']) if opt.single_cls else (int(
data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
# 加载预训练的模型参数
if pretrained:
# 下载与训练数据
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get('anchors'):
ckpt['model'].yaml['anchors'] = round(
hyp['anchors']) # force autoanchor
# 载入输入的配置或者是加载的pretrained的配置，ch=3是输入channel
model = Model(opt.cfg or ckpt['model'].yaml, ch=3,
nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [
] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
# 只加载在预训练的模型和当前模型中都有的组件的参数，这要求与训练的模型和当前模型的shape要相等
# intersect_dicts的左右就是将与训练的模型参数和当前的模型参数进行比较，取shape一致的那些参数（shape不一样的是没法运用在当前模型的）
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
# 非严格模式加载参数
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
# 冻结某几层，finetune可以用
freeze = [
'',
] # parameter names to freeze (full or partial)
if any(freeze):
for k, v in model.named_parameters():
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
# 将batch_size和64比较，当64不是batch_size的整数倍的时候，权重做相应的调整
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
v.requires_grad = True
if '.bias' in k:
pg2.append(v) # biases
elif '.weight' in k and '.bn' not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
# 添加卷积权重参数和biases参数，其中biases不需要权重衰减，这里的params只能是这个名字，执行完add_param_group后optimizer的数据就是一个list中有三个值
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
# 权重衰减策略
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp[
'lrf']) + hyp['lrf'] # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
shutil.copytree(wdir, wdir.parent /
f'weights_backup_epoch{start_epoch - 1}'
) # save previous weights
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
# 得到处理（就是让图片大小能被网格大小整除）后的图像和测试图像的大小
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# DP mode
# 分布式（单机多GPU）
if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
# 同步batchnorm，如果不同步的话每个GPU上的batchnorm都会使用当前GPU上数据的方差和均值，那几个GPU虽然训练的是同一个batch的数据，值却是不一样的
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# Exponential moving average 指数移动平均
# 给予近期数据更高的权重，就是说对于参数，我们给予最近的几次的参数更高的权重，其假设就是最近几次的参数是在最优处抖动，所以最近几次的参数权重就给高点
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
# 分布式（多机器多GPU）
if cuda and rank != -1:
# local_rank指定的是当前进程使用的是哪块GPU，local_rank表示的就是GPU序号
model = DDP(model,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# Process 0
if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates
testloader = create_dataloader(test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images
and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers)[0] # testloader
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
plot_labels(labels, save_dir=log_dir)
if tb_writer:
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
tb_writer.add_histogram('classes', c, 0)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device) # attach class weights
model.names = names
# Start training
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * nb),
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
logger.info('Image sizes %g train, %g test\n'
'Using %g dataloader workers\nLogging results to %s\n'
'Starting training for %g epochs...' %
(imgsz, imgsz_test, dataloader.num_workers, log_dir, epochs))
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if rank in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (
1 - maps)**2 # class weights
iw = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=cw) # image weights
dataset.indices = random.choices(
range(dataset.n), weights=iw,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = (torch.tensor(dataset.indices)
if rank == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls',
'total', 'targets', 'img_size'))
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device),
model) # loss scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if ni < 3:
f = str(log_dir / ('train_batch%g.jpg' % ni)) # filename
result = plot_images(images=imgs,
targets=targets,
paths=paths,
fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f,
result,
dataformats='HWC',
global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema:
ema.update_attr(
model,
include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir,
plots=epoch == 0 or final_epoch) # plot first and last
# Write
with open(results_file, 'a') as f:
f.write(
s + '%10.4g' * 7 % results +
'\n') # P, R, [email protected], [email protected], val_loss(box, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' %
(results_file, opt.bucket, opt.name))
# Tensorboard
if tb_writer:
tags = [
'train/box_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
tb_writer.add_scalar(tag, x, epoch)
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
f.read(),
'model':
ema.ema,
'optimizer':
None if final_epoch else optimizer.state_dict()
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
n = opt.name if opt.name.isnumeric() else ''
fresults, flast, fbest = log_dir / f'results{n}.txt', wdir / f'last{n}.pt', wdir / f'best{n}.pt'
for f1, f2 in zip([wdir / 'last.pt', wdir / 'best.pt', results_file],
[flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
if str(f2).endswith('.pt'): # is *.pt
strip_optimizer(f2) # strip optimizer
os.system(
'gsutil cp %s gs://%s/weights' %
(f2, opt.bucket)) if opt.bucket else None # upload
# Finish
if not opt.evolve:
plot_results(save_dir=log_dir) # save as results.png
logger.info('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
return results
from utils.google_utils import attempt_download attempt_download(f'yolov5m.pt')
def __detect(image, opt, save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt['output'], opt['source'], opt['weights'], opt['view_img'], opt['save_txt'], opt['img_size']
webcam = False
# Initialize
device = torch_utils.select_device(opt['device'])
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# half = device.type != 'cpu' # half precision only supported on CUDA
half = False
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights, map_location=device)['model'].float() # load to FP32
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
imgsz = check_img_size(imgsz, s=model.model[-1].stride.max()) # check img_size
# Set Dataloader
dataset = LoadInMemoryImage(image, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
# Run inference
res = []
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
for img, im0s in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment = opt['augment'])[0]
# Apply NMS
pred = non_max_suppression(pred, opt['conf_thres'], opt['iou_thres'], classes=opt['classes'], agnostic=opt['agnostic_nms'])
t2 = torch_utils.time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
im0 = im0s
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Write results
detections = []
for *xyxy, conf, cls in det:
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
detections.append({
'class': cls.item(),
'x': xywh[0],
'y': xywh[1],
'w': xywh[2],
'h': xywh[3],
'confidence': conf
})
res.append(detections)
return res
def train(hyp, opt, device):
save_dir, epochs, batch_size, total_batch_size, weights, rank = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
do_semi = opt.do_semi
# Directories
wdir = save_dir / 'weights'
wdir.mkdir(parents=True, exist_ok=True)
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = save_dir / 'results.txt'
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
plots = not opt.evolve #create plots
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader)
nc = 1 if opt.single_cls else int(data_dict['nc']) #number of classes
names = ['item'] if opt.single_cls and len(
data_dict['names']) != 1 else data_dict['names']
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data)
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) #load checkpoint
model = Model(opt.cfg or ckpt['model'].yaml,
ch=3,
nc=nc,
anchors=hyp.get('anchors')).to(device) #create
exclude = [
'anchor'
] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [
] #exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) #intersect
model.load_state_dict(state_dict, strict=False) #load
else:
model = Model(opt.cfg, ch=3, nc=nc,
anchors=hyp.get('anchors')).to(device)
with torch_distributed_zero_first(rank):
check_dataset(data_dict) #check
train_path = data_dict['train']
test_path = data_dict['val']
# Optimizer
nbs = 64
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply dacay
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust betal to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
del pg0, pg1, pg2
if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp[
'lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs)
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# EMA
ema = ModelEMA(model) if rank in [-1, 0] else None
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Results
if ckpt.get('training_results') is not None:
results_file.write_text(
ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weight, epochs)
if epochs < start_epoch:
epochs += ckpt['epoch']
del ckpt, state_dict
# Image sizes
gs = max(int(model.stride.max()), 32) # grid size (max stride)
nl = model.model[
-1].nl # number of detection layer (used for scaling hyp['obj])
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(
device)
# Trainloader
if do_semi:
dataloader, dataset, unlabeldataloader = create_dataloader(
train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
do_semi=opt.do_semi)
else:
dataloader, dataset = create_dataloader(
train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
do_semi=opt.do_semi)
# Train teacher model
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# process 0
if rank in [-1, 0]:
testloader = create_dataloader(
test_path,
imgsz_test,
batch_size * 2,
gs,
opt, # testloader
hyp=hyp,
cache=opt.cache_images and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers,
pad=0.5,
prefix=colorstr('val: '),
do_semi=False)[0]
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
model.half().float() # pre-reduce anchor precision
# DDP mode
if cuda and rank != 1:
model = DDP(model,
device_ids=[opt.local_rank],
output_device=opt.local_rank,
find_unused_parameters=any(
isinstance(layer, nn.MultiheadAttention)
for layer in model.modules()))
# Model parameters
hyp['box'] *= 3. / nl # scale to layers
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640)**2 * 3. / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(
dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Train teacher model --> burn in
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * nb),
1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
compute_loss = ComputeLoss(model) # init loss class
burnin_epochs = epochs / 2
# burn in
for epoch in range(start_epoch,
burnin_epochs): # epoch-------------------------
model.train()
nb = len(dataloader)
mloss = torch.zeros(4, device=device) # mean loss
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb)
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar:
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warm up
if ni <= [0, nw]:
xi = [0, nw]
accumulate = max(
1, np.interp(ni, xi, [1, nbs / total_batch_size].round()))
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_item = compute_loss(
pred, targets.to(device)) # loss scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between device in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# print
if rank in [-1, 0]:
mloss = (mloss * i + loss_item) / (i + 1) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
ema.update_attr(model,
include=[
'yaml', 'nc', 'hyp', 'gr', 'names', 'stride',
'class_weights'
])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(data_dict,
batch_size=batch_size * 2,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
verbose=nc < 50
and final_epoch,
plots=plots and final_epoch,
compute_loss=compute_loss)
fi = fitness(np.array(results).reshape(
1, -1)) # weighted combination of [P, R, mAP@50, [email protected]]
if fi > best_fitness:
best_fitness = fi
if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
results_file.read_text(),
'model':
deepcopy(model.module if is_parallel(model) else model).half(),
'ema':
deepcopy(ema.ema).half(),
'updates':
ema.updates,
'optimizer':
optimizer.state_dict()
}
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------
# end warm up
# get persudo label
# STAC
# first apply weak augmentation on unlabeled dataset then use teacher net to predict the persudo labels
# Then apply strong augmentation on unlabeled dataset, use student net to get the logists and compute the unlabeled loss.
model.eval()
img = []
target = []
Path = []
imgsz = opt.img_size
for idx, batch in tqdm(enumerate(unlabeldataloader),
total=len(unlabeldataloader)):
imgs0, _, path, _ = batch # from uint8 to float16
with torch.no_grad():
pred = model(imgs0.to(device, non_blocking=True).float() /
255.0)[0]
gn = torch.tensor(imgs0.shape)[[3, 2, 3, 2]]
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
for index, pre in enumerate(pred):
predict_number = len(pre)
if predict_number == 0:
continue
Class = pre[:, 5].view(predict_number, 1).cpu()
XYWH = (xyxy2xywh(pre[:, :4])).cpu()
XYWH /= gn
pre = torch.cat((torch.zeros(predict_number, 1), Class, XYWH),
dim=1)
img.append(imgs0[index])
target.append(pre)
Path.append(path[index])
unlabeldataset = semiDataset(img, target, Path)
del img, targets, Path
model.train()
def train(hyp, logger, work_dir, device):
epochs = opt.epochs
batch_size = opt.batch_size
total_batch_size = opt.total_batch_size
weights = opt.weights
rank = opt.local_rank
# Configure
init_seeds(1)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
#train_path = data_dict['train']
#test_path = data_dict['val']
train_path = os.path.join(data_dict['convertor_path'], 'images',
'train2017')
test_path = os.path.join(data_dict['convertor_path'], 'images', 'val2017')
nc, names = (1, ['item']) if opt.single_cls else (int(
len(data_dict['names'])), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Create model
model = Model(opt.cfg, nc=nc).to(device)
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# Optimizer
nbs = 64 # nominal batch size
# default DDP implementation is slow for accumulation according to: https://pytorch.org/docs/stable/notes/ddp.html
# all-reduce operation is carried out during loss.backward().
# Thus, there would be redundant all-reduce communications in a accumulation procedure,
# which means, the result is still right but the training speed gets slower.
# TODO: If acceleration is needed, there is an implementation of allreduce_post_accumulation
# in https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/run_pretraining.py
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
if v.requires_grad:
if '.bias' in k:
pg2.append(v) # biases
elif '.weight' in k and '.bn' not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
if hyp['optimizer'] == 'adam': # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Load Model
google_utils.attempt_download(weights)
start_epoch, best_fitness = 0, 0.0
# 加载自己的模型
if not weights.endswith('.pt'):
ckpt = torch.load(weights, map_location=device).float()
model.load_state_dict(ckpt.state_dict(), strict=True)
logger.info(f'load myself ckpt: {weights}')
if weights.endswith('.pt'): # pytorch format
ckpt = torch.load(weights, map_location=device) # load checkpoint
# load model
try:
exclude = ['anchor'] # exclude keys
ckpt['model'] = {
k: v
for k, v in ckpt['model'].float().state_dict().items()
if k in model.state_dict() and not any(x in k for x in exclude)
and model.state_dict()[k].shape == v.shape
}
model.load_state_dict(ckpt['model'], strict=False)
print('Transferred %g/%g items from %s' %
(len(ckpt['model']), len(model.state_dict()), weights))
except KeyError as e:
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \
"Please delete or update %s and try again, or use --weights '' to train from scratch." \
% (weights, opt.cfg, weights, weights)
raise KeyError(s) from e
# load optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# load results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# epochs
start_epoch = ckpt['epoch'] + 1
if epochs < start_epoch:
print(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt
# Mixed precision training https://github.com/NVIDIA/apex
if mixed_precision:
model, optimizer = amp.initialize(model,
optimizer,
opt_level='O1',
verbosity=0)
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
lf = lambda x: ((
(1 + math.cos(x * math.pi / epochs)) / 2)**1.0) * 0.8 + 0.2 # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# https://discuss.pytorch.org/t/a-problem-occured-when-resuming-an-optimizer/28822
# plot_lr_scheduler(optimizer, scheduler, epochs)
# DP mode
if device.type != 'cpu' and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and device.type != 'cpu' and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# Exponential moving average
ema = torch_utils.ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if device.type != 'cpu' and rank != -1:
model = DDP(model, device_ids=[rank], output_device=rank)
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
local_rank=rank,
world_size=opt.world_size)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# Testloader
if rank in [-1, 0]:
# local_rank is set to -1. Because only the first process is expected to do evaluation.
testloader = create_dataloader(test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images,
rect=True,
local_rank=-1,
world_size=opt.world_size)[0]
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device) # attach class weights
model.names = names
# Class frequency
if rank in [-1, 0]:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1.
# model._initialize_biases(cf.to(device))
# Check anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# save anchors
m = model.module.model[-1] if hasattr(model,
'module') else model.model[-1]
anchors = []
for i in range(3):
for j in range(3):
anchor = m.anchor_grid[i, 0, j, 0,
0].cpu().detach().numpy().tolist()
anchors.append(anchor)
with open(os.path.join(work_dir, 'anchors.txt'), 'w') as f:
for anchor in anchors:
f.write(f'{anchor[0]},{anchor[1]}\n')
# Start training
t0 = time.time()
nw = max(3 * nb,
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
maps = np.zeros(nc) # mAP per class
results = (
0, 0, 0, 0, 0, 0, 0
) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
scheduler.last_epoch = start_epoch - 1 # do not move
if rank in [0, -1]:
logger.info('Image sizes %g train, %g test' % (imgsz, imgsz_test))
logger.info('Using %g dataloader workers' % dataloader.num_workers)
logger.info('Starting training for %g epochs...' % epochs)
# torch.autograd.set_detect_anomaly(True)
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
train_time_start = time.time()
logger.info('')
logger.info('epoch: {epoch} lr: {lr}'.format(
epoch=epoch, lr=optimizer.param_groups[0]['lr']))
model.train()
# Update image weights (optional)
# When in DDP mode, the generated indices will be broadcasted to synchronize dataset.
if dataset.image_weights:
# Generate indices.
if rank in [-1, 0]:
w = model.class_weights.cpu().numpy() * (
1 - maps)**2 # class weights
image_weights = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=w)
dataset.indices = random.choices(
range(dataset.n), weights=image_weights,
k=dataset.n) # rand weighted idx
# Broadcast.
if rank != -1:
indices = torch.zeros([dataset.n], dtype=torch.int)
if rank == 0:
indices[:] = torch.from_tensor(dataset.indices,
dtype=torch.int)
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
'''
pbar = enumerate(dataloader)
if rank in [-1, 0]:
logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls', 'total', 'targets', 'img_size'))
pbar = tqdm(pbar, total=nb) # progress bar
'''
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in enumerate(
dataloader
): # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0 - 255 to 0.0 - 1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(
ni, xi,
[0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi,
[0.9, hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
pred = model(imgs)
# Loss
loss, loss_items = compute_loss(pred, targets.to(device),
model) # scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
if not torch.isfinite(loss):
logger.info('WARNING: non-finite loss, ending training ',
loss_items)
return results
# Backward
if mixed_precision:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# Optimize
if ni % accumulate == 0:
optimizer.step()
optimizer.zero_grad()
if ema is not None:
ema.update(model)
if i % 200 == 0:
logger.info(
'[Epoch:{epoch}/{epochs} iter:{iter}] loss:{loss}'.format(
epoch=epoch,
epochs=epochs - 1,
iter=i,
loss=loss.item()))
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
scheduler.step()
train_time_end = time.time()
logger.info('train time: {train_time}s'.format(
train_time=int(train_time_end - train_time_start)))
# Only the first process in DDP mode is allowed to log or save checkpoints.
if rank in [-1, 0]:
# mAP
if ema is not None:
ema.update_attr(
model,
include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride'])
final_epoch = epoch + 1 == epochs
if (epoch % data_dict['eval_interval'] == 0
and epoch != 0) or final_epoch: # Calculate mAP
results, maps, times = test.test(
data_dict,
batch_size=total_batch_size,
imgsz=imgsz_test,
save_json=final_epoch
and opt.data.endswith(os.sep + 'coco.yaml'),
model=ema.ema.module
if hasattr(ema.ema, 'module') else ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=work_dir)
map50, map = results[2], results[3]
logger.info(f'eval: [email protected]: {map50} [email protected]:.95: {map}')
# 保存模型
ckpt = ema.ema.module if hasattr(ema.ema,
'module') else ema.ema
torch.save(
ckpt,
os.path.join(work_dir,
'epoch_{epoch}.pth'.format(epoch=epoch)))
# end epoch ----------------------------------------------------------------------------------------------------
# end training
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
return results
def train(hyp):
print(f"Hyperparameters {hyp}")
log_dir = tb_writer.log_dir if tb_writer else "runs/evolution" # run directory
wdir = str(Path(log_dir) / "weights") + os.sep # weights directory
os.makedirs(wdir, exist_ok=True)
last = wdir + "last.pt"
best = wdir + "best.pt"
results_file = log_dir + os.sep + "results.txt"
# Save run settings
with open(Path(log_dir) / "hyp.yaml", "w") as f:
yaml.dump(hyp, f, sort_keys=False)
with open(Path(log_dir) / "opt.yaml", "w") as f:
yaml.dump(vars(opt), f, sort_keys=False)
epochs = opt.epochs # 300
batch_size = opt.batch_size # 64
weights = opt.weights # initial training weights
# Configure
init_seeds(1)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
train_path = data_dict["train"]
test_path = data_dict["val"]
nc, names = (
(1, ["item"]) if opt.single_cls else
(int(data_dict["nc"]), data_dict["names"])) # number classes, names
assert len(names) == nc, "%g names found for nc=%g dataset in %s" % (
len(names),
nc,
opt.data,
) # check
# Remove previous results
for f in glob.glob("*_batch*.jpg") + glob.glob(results_file):
os.remove(f)
# Create model
model = Model(opt.cfg, nc=nc).to(device)
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size),
1) # accumulate loss before optimizing
hyp["weight_decay"] *= batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
if v.requires_grad:
if ".bias" in k:
pg2.append(v) # biases
elif ".weight" in k and ".bn" not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
if (
hyp["optimizer"] == "adam"
): # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
optimizer = optim.Adam(pg0,
lr=hyp["lr0"],
betas=(hyp["momentum"],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp["lr0"],
momentum=hyp["momentum"],
nesterov=True)
optimizer.add_param_group({
"params": pg1,
"weight_decay": hyp["weight_decay"]
}) # add pg1 with weight_decay
optimizer.add_param_group({"params": pg2}) # add pg2 (biases)
print("Optimizer groups: %g .bias, %g conv.weight, %g other" %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
lf = lambda x: ((
(1 + math.cos(x * math.pi / epochs)) / 2)**1.0) * 0.9 + 0.1 # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs, save_dir=log_dir)
# Load Model
google_utils.attempt_download(weights)
start_epoch, best_fitness = 0, 0.0
if weights.endswith(".pt"): # pytorch format
ckpt = torch.load(weights, map_location=device) # load checkpoint
# load model
try:
ckpt["model"] = {
k: v
for k, v in ckpt["model"].float().state_dict().items()
if model.state_dict()[k].shape == v.shape
} # to FP32, filter
model.load_state_dict(ckpt["model"], strict=False)
except KeyError as e:
s = (
"%s is not compatible with %s. This may be due to model differences or %s may be out of date. "
"Please delete or update %s and try again, or use --weights '' to train from scratch."
% (opt.weights, opt.cfg, opt.weights, opt.weights))
raise KeyError(s) from e
# load optimizer
if ckpt["optimizer"] is not None:
optimizer.load_state_dict(ckpt["optimizer"])
best_fitness = ckpt["best_fitness"]
# load results
if ckpt.get("training_results") is not None:
with open(results_file, "w") as file:
file.write(ckpt["training_results"]) # write results.txt
# epochs
start_epoch = ckpt["epoch"] + 1
if epochs < start_epoch:
print(
"%s has been trained for %g epochs. Fine-tuning for %g additional epochs."
% (opt.weights, ckpt["epoch"], epochs))
epochs += ckpt["epoch"] # finetune additional epochs
del ckpt
# Mixed precision training https://github.com/NVIDIA/apex
if mixed_precision:
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O1",
verbosity=0)
# Distributed training
if device.type != "cpu" and torch.cuda.device_count(
) > 1 and dist.is_available():
dist.init_process_group(
backend="nccl", # distributed backend
init_method="tcp://127.0.0.1:9999", # init method
world_size=1, # number of nodes
rank=0,
) # node rank
# model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) # requires world_size > 1
model = torch.nn.parallel.DistributedDataParallel(model)
# Trainloader
dataloader, dataset = create_dataloader(
train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, "Label class %g exceeds nc=%g in %s. Correct your labels or your model." % (
mlc,
nc,
opt.cfg,
)
# Testloader
testloader = create_dataloader(
test_path,
imgsz_test,
batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images,
rect=True,
)[0]
# Model parameters
hyp["cls"] *= nc / 80.0 # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device) # attach class weights
model.names = names
# Class frequency
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1.
# model._initialize_biases(cf.to(device))
plot_labels(labels, save_dir=log_dir)
if tb_writer:
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
tb_writer.add_histogram("classes", c, 0)
# Check anchors
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz)
# Exponential moving average
ema = torch_utils.ModelEMA(model)
# Start training
t0 = time.time()
nw = max(3 * nb,
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
maps = np.zeros(nc) # mAP per class
results = (
0,
0,
0,
0,
0,
0,
0,
) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
scheduler.last_epoch = start_epoch - 1 # do not move
print("Image sizes %g train, %g test" % (imgsz, imgsz_test))
print("Using %g dataloader workers" % dataloader.num_workers)
print("Starting training for %g epochs..." % epochs)
# torch.autograd.set_detect_anomaly(True)
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if dataset.image_weights:
w = model.class_weights.cpu().numpy() * (1 -
maps)**2 # class weights
image_weights = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=w)
dataset.indices = random.choices(range(dataset.n),
weights=image_weights,
k=dataset.n) # rand weighted idx
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
print(("\n" + "%10s" * 8) % ("Epoch", "gpu_mem", "GIoU", "obj", "cls",
"total", "targets", "img_size"))
pbar = tqdm(enumerate(dataloader), total=nb) # progress bar
for (
i,
(imgs, targets, paths, _),
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = (imgs.to(device, non_blocking=True).float() / 255.0
) # uint8 to float32, 0 - 255 to 0.0 - 1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x["lr"] = np.interp(
ni, xi,
[0.1 if j == 2 else 0.0, x["initial_lr"] * lf(epoch)])
if "momentum" in x:
x["momentum"] = np.interp(ni, xi,
[0.9, hyp["momentum"]])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode="bilinear",
align_corners=False)
# Forward
pred = model(imgs)
# Loss
loss, loss_items = compute_loss(pred, targets.to(device), model)
if not torch.isfinite(loss):
print("WARNING: non-finite loss, ending training ", loss_items)
return results
# Backward
if mixed_precision:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# Optimize
if ni % accumulate == 0:
optimizer.step()
optimizer.zero_grad()
ema.update(model)
# Print
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = "%.3gG" % (torch.cuda.memory_cached() /
1e9 if torch.cuda.is_available() else 0) # (GB)
s = ("%10s" * 2 + "%10.4g" * 6) % (
"%g/%g" % (epoch, epochs - 1),
mem,
*mloss,
targets.shape[0],
imgs.shape[-1],
)
pbar.set_description(s)
# Plot
if ni < 3:
f = str(Path(log_dir) / ("train_batch%g.jpg" % ni)) # filename
result = plot_images(images=imgs,
targets=targets,
paths=paths,
fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f,
result,
dataformats="HWC",
global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
scheduler.step()
# mAP
ema.update_attr(model,
include=["md", "nc", "hyp", "gr", "names", "stride"])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(
opt.data,
batch_size=batch_size,
imgsz=imgsz_test,
save_json=final_epoch
and opt.data.endswith(os.sep + "coco.yaml"),
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir,
)
# Write
with open(results_file, "a") as f:
f.write(s + "%10.4g" * 7 % results +
"\n") # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
if len(opt.name) and opt.bucket:
os.system("gsutil cp %s gs://%s/results/results%s.txt" %
(results_file, opt.bucket, opt.name))
# Tensorboard
if tb_writer:
tags = [
"train/giou_loss",
"train/obj_loss",
"train/cls_loss",
"metrics/precision",
"metrics/recall",
"metrics/mAP_0.5",
"metrics/mAP_0.5:0.95",
"val/giou_loss",
"val/obj_loss",
"val/cls_loss",
]
for x, tag in zip(list(mloss[:-1]) + list(results), tags):
tb_writer.add_scalar(tag, x, epoch)
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, "r") as f: # create checkpoint
ckpt = {
"epoch": epoch,
"best_fitness": best_fitness,
"training_results": f.read(),
"model": ema.ema,
"optimizer":
None if final_epoch else optimizer.state_dict(),
}
# Save last, best and delete
torch.save(ckpt, last)
if (best_fitness == fi) and not final_epoch:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
# Strip optimizers
n = ("_" if len(opt.name) and not opt.name.isnumeric() else "") + opt.name
fresults, flast, fbest = "results%s.txt" % n, wdir + "last%s.pt" % n, wdir + "best%s.pt" % n
for f1, f2 in zip([wdir + "last.pt", wdir + "best.pt", "results.txt"],
[flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
ispt = f2.endswith(".pt") # is *.pt
strip_optimizer(f2) if ispt else None # strip optimizer
os.system(
"gsutil cp %s gs://%s/weights" %
(f2, opt.bucket)) if opt.bucket and ispt else None # upload
# Finish
if not opt.evolve:
plot_results(save_dir=log_dir) # save as results.png
print("%g epochs completed in %.3f hours.\n" % (epoch - start_epoch + 1,
(time.time() - t0) / 3600))
dist.destroy_process_group(
) if device.type != "cpu" and torch.cuda.device_count() > 1 else None
torch.cuda.empty_cache()
return results
def train(hyp, opt, device):
logger.info(
colorstr('hyperparameters: ') + ', '.join(f'{k}={v}'
for k, v in hyp.items()))
save_dir, epochs, batch_size, total_batch_size, weights, rank = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Directories
wdir = save_dir / 'weights'
wdir.mkdir(parents=True, exist_ok=True) # make dir
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = save_dir / 'results.txt'
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
plots = not opt.evolve # create plots
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.SafeLoader) # data dict
is_coco = opt.data.endswith('coco.yaml')
# Logging- Doing this before checking the dataset. Might update data_dict
opt.hyp = hyp # add hyperparameters
nc = int(data_dict['nc']) # number of classes
names = data_dict['names'] # class names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
if pretrained:
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
model = Model(opt.cfg or ckpt['model'].yaml,
ch=3,
nc=nc,
anchors=hyp.get('anchors')).to(device) # create
exclude = [
'anchor'
] if (opt.cfg or hyp.get('anchors')) and not opt.resume else [
] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc,
anchors=hyp.get('anchors')).to(device) # create
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
# accumulate loss before optimizing
accumulate = max(round(nbs / total_batch_size), 1)
hyp['weight_decay'] *= total_batch_size * \
accumulate / nbs # scale weight_decay
logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
# add pg1 with weight_decay
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
})
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
if opt.linear_lr:
def lf(x): return (1 - x / (epochs - 1)) * \
(1.0 - hyp['lrf']) + hyp['lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA
ema = ModelEMA(model) if rank == -1 else None
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Results
if ckpt.get('training_results') is not None:
results_file.write_text(
ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = max(int(model.stride.max()), 32) # grid size (max stride)
# number of detection layers (used for scaling hyp['obj'])
nl = model.model[-1].nl
# verify imgsz are gs-multiples
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size]
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
world_size=opt.world_size,
workers=opt.workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '))
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# Process 0
if rank == -1:
testloader = create_dataloader(
test_path,
imgsz_test,
batch_size * 2,
gs,
opt, # testloader
hyp=hyp,
cache=opt.cache_images and not opt.notest,
rect=True,
world_size=opt.world_size,
workers=opt.workers,
pad=0.5,
prefix=colorstr('val: '))[0]
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
if plots:
plot_labels(labels, names, save_dir)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
model.half().float() # pre-reduce anchor precision
# Model parameters
hyp['box'] *= 3. / nl # scale to layers
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
# scale to image size and layers
hyp['obj'] *= (imgsz / 640)**2 * 3. / nl
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(
dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Start training
t0 = time.time()
# number of warmup iterations, max(3 epochs, 1k iterations)
nw = max(round(hyp['warmup_epochs'] * nb), 1000)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
# P, R, [email protected], [email protected], val_loss(box, obj, cls)
results = (0, 0, 0, 0, 0, 0, 0)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
compute_loss = ComputeLoss(model) # init loss class
logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
f'Using {dataloader.num_workers} dataloader workers\n'
f'Logging results to {save_dir}\n'
f'Starting training for {epochs} epochs...')
# epoch ------------------------------------------------------------------
for epoch in range(start_epoch, epochs):
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
cw = model.class_weights.cpu().numpy() * (
1 - maps)**2 / nc # class weights
iw = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n),
weights=iw,
k=dataset.n) # rand weighted idx
mloss = torch.zeros(4, device=device) # mean losses
pbar = enumerate(dataloader)
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls',
'total', 'labels', 'img_size'))
if rank == -1:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
# batch -------------------------------------------------------------
for i, (imgs, targets, paths, _) in pbar:
# number integrated batches (since train start)
ni = i + nb * epoch
imgs = imgs.to(device, non_blocking=True).float() / \
255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
# new shape (stretched to gs-multiple)
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device)) # loss scaled by batch_size
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank == -1:
mloss = (mloss * i + loss_items) / \
(i + 1) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if plots and ni < 3:
f = save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images,
args=(imgs, targets, paths, f),
daemon=True).start()
# end batch ------------------------------------------------------------------------------------------------
# end epoch ----------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# single-GPU
if rank == -1:
# mAP
ema.update_attr(model,
include=[
'yaml', 'nc', 'hyp', 'gr', 'names', 'stride',
'class_weights'
])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(data_dict,
batch_size=batch_size * 2,
imgsz=imgsz_test,
model=ema.ema,
dataloader=testloader,
save_dir=save_dir,
verbose=nc < 50
and final_epoch,
plots=plots and final_epoch,
compute_loss=compute_loss,
is_coco=is_coco)
# Write
with open(results_file, 'a') as f:
# append metrics, val_loss
f.write(s + '%10.4g' * 7 % results + '\n')
# Log
tags = [
'train/box_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/box_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
# Update best mAP
# weighted combination of [P, R, [email protected], [email protected]]
fi = fitness(np.array(results).reshape(1, -1))
if fi > best_fitness:
best_fitness = fi
# Save model
if (not opt.nosave) or (final_epoch and not opt.evolve): # if save
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'training_results': results_file.read_text(),
'model': deepcopy(model).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict()
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank == -1:
# Plots
if plots:
plot_results(save_dir=save_dir) # save as results.png
# Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for m in (last,
best) if best.exists() else (last): # speed, mAP tests
results, _, _ = test.test(opt.data,
batch_size=batch_size * 2,
imgsz=imgsz_test,
conf_thres=0.001,
iou_thres=0.7,
model=attempt_load(m, device).half(),
dataloader=testloader,
save_dir=save_dir,
save_json=True,
plots=False,
is_coco=is_coco)
# Strip optimizers
for f in last, best:
if f.exists():
strip_optimizer(f) # strip optimizers
else:
dist.destroy_process_group()
torch.cuda.empty_cache()
return results
def train(hyp, opt, device, tb_writer=None):
logger.info(f'Hyperparameters {hyp}')
log_dir = Path(tb_writer.log_dir) if tb_writer else Path(
opt.logdir) / 'evolve' # logging directory
wdir = log_dir / 'weights' # weights directory
os.makedirs(wdir, exist_ok=True)
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = str(log_dir / 'results.txt')
epochs, batch_size, total_batch_size, weights, rank = \
opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Save run settings
with open(log_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(log_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc, names = (1, ['item']) if opt.single_cls else (int(
data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
# added by jiangrong
if not opt.resume:
ckpt['epoch'] = -1
if opt.nas:
model = NasModel(opt.cfg,
ch=3,
nc=nc,
nas=opt.nas,
nas_stage=opt.nas_stage).to(device) # create
else:
model = Model(opt.cfg or ckpt['model'].yaml, ch=3,
nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg else [] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
if opt.nas:
model = NasModel(opt.cfg,
ch=3,
nc=nc,
nas=opt.nas,
nas_stage=opt.nas_stage).to(device) # create
if opt.nas_stage == 3:
# TODO, Remapping with BN Statistics on Width-level
model.re_organize_middle_weights()
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
if opt.nas and opt.nas_stage > 0:
from models.experimental import attempt_load
"""
P R [email protected]
0.535 0.835 0.742
python test.py \
--weights /workspace/yolov5-v3/yolov5/runs/exp122/weights/best.pt \
--data ./data/baiguang.yaml \
--device 1 \
--conf-thres 0.2
"""
teacher_model = attempt_load(
"/workspace/yolov5-v3/yolov5/runs/exp259/weights/best.pt",
map_location='cuda:1')
teacher_model.eval()
# Freeze
freeze = [
'',
] # parameter names to freeze (full or partial)
if any(freeze):
for k, v in model.named_parameters():
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
v.requires_grad = True
if '.bias' in k:
pg2.append(v) # biases
elif '.weight' in k and '.bn' not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
if opt.adam:
optimizer = optim.Adam(pg0,
lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = lambda x: ((
(1 + math.cos(x * math.pi / epochs)) / 2)**1.0) * 0.8 + 0.2 # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None and not opt.nas > 0:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
shutil.copytree(wdir, wdir.parent /
f'weights_backup_epoch{start_epoch - 1}'
) # save previous weights
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# DP mode
# TheModel = model
if cuda and rank == -1 and torch.cuda.device_count() > 1 and not (
opt.nas and opt.nas_stage > 0):
# https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html
# >>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
# >>> output = net(input_var) # input_var can be on any device, including CPU
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# Exponential moving average
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if cuda and rank != -1:
model = DDP(model,
device_ids=[opt.local_rank],
output_device=(opt.local_rank))
# Trainloader
dataloader, dataset = create_dataloader(train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
# Testloader
if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates
testloader = create_dataloader(
test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers)[0] # only runs on process 0
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device) # attach class weights
model.names = names
# Class frequency
if rank in [-1, 0]:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1.
# model._initialize_biases(cf.to(device))
plot_labels(labels, save_dir=log_dir)
if tb_writer:
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
tb_writer.add_histogram('classes', c, 0)
# Check anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# Start training
t0 = time.time()
nw = max(3 * nb,
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (
0, 0, 0, 0, 0, 0, 0
) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
scheduler.last_epoch = start_epoch - 1 # do not move
# scaler = amp.GradScaler(enabled=cuda)
logger.info('Image sizes %g train, %g test' % (imgsz, imgsz_test))
logger.info('Using %g dataloader workers' % dataloader.num_workers)
logger.info('Starting training for %g epochs...' % epochs)
# torch.autograd.set_detect_anomaly(True)
plot_csum = 0
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if dataset.image_weights:
# Generate indices
if rank in [-1, 0]:
w = model.class_weights.cpu().numpy() * (
1 - maps)**2 # class weights
image_weights = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=w)
dataset.indices = random.choices(
range(dataset.n), weights=image_weights,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = torch.zeros([dataset.n], dtype=torch.int)
if rank == 0:
indices[:] = torch.tensor(dataset.indices, dtype=torch.int)
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls',
'total', 'targets', 'img_size'))
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
# print(type(targets), targets.size()) # [[_,classid(start from 0), x,y,w,h (0-1)]]
# print('---> targets: ', targets)
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(
ni, xi,
[0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi,
[0.9, hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
###### jiangrong, turn off mixed precision ##########
# with amp.autocast(enabled=cuda):
if 1 == 1:
pred = model(imgs) # forward, format x(bs,3,20,20,80+1+4)
loss, loss_items = compute_loss(
pred, targets.to(device),
model) # loss scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
# z= []
# for i in range(TheModel._modules['model'][-1].nl):
# bs, _, ny, nx, _ = pred[i].shape
# if TheModel._modules['model'][-1].grid[i].shape[2:4] != pred[i].shape[2:4]:
# TheModel._modules['model'][-1].grid[i] = TheModel._modules['model'][-1]._make_grid(nx, ny).to(pred[i].device)
# y = pred[i].sigmoid()
# y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + TheModel._modules['model'][-1].grid[i].to(pred[i].device)) * TheModel._modules['model'][-1].stride[i] # xy
# y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * TheModel._modules['model'][-1].anchor_grid[i] # wh
# z.append(y.view(bs, -1, TheModel._modules['model'][-1].no))
# inf_out = torch.cat(z, 1)
# teacher_pred = non_max_suppression(inf_out, conf_thres=0.2, iou_thres=0.6, merge=False)
# assert len(teacher_pred) == imgs.size()[0]
# for i, (det, plot_img) in enumerate(zip(teacher_pred, imgs.detach().cpu().numpy())):
# plot_img = np.transpose(plot_img, (1,2,0))
# plot_img = np.uint8(plot_img * 255.0)
# plot_csum += 1
# cv2.imwrite('./tmp/{}.jpg'.format(plot_csum), plot_img)
# plot_img = cv2.imread('./tmp/{}.jpg'.format(plot_csum))
# for tgt in targets.detach().cpu().numpy():
# _, tgt_class_id, c_x, c_y, c_w, c_h = tgt
# c_x, c_y, c_w, c_h = float(c_x), float(c_y), float(c_w), float(c_h)
# c_x, c_y, c_w, c_h = c_x * plot_img.shape[1], c_y * plot_img.shape[0], c_w * plot_img.shape[1], c_h * plot_img.shape[0]
# cv2.rectangle(plot_img, (int(c_x - c_w / 2), int(c_y - c_h / 2)), (int(c_x + c_w / 2), int(c_y + c_h / 2)), (0,0,255), 2)
# print('===> ', int(c_x - c_w / 2), int(c_y - c_h / 2), int(c_x + c_w / 2), int(c_y + c_h / 2), tgt_class_id)
# if det is not None:
# det = det.detach().cpu().numpy()
# for each_b in det:
# pass
# cv2.rectangle(plot_img, (int(each_b[0]), int(each_b[1])), (int(each_b[2]), int(each_b[3])), (255,0,0), 2)
# print('---> ', int(each_b[0]), int(each_b[1]), int(each_b[2]), int(each_b[3]), float(each_b[4]), int(each_b[5]))
# cv2.imwrite('./tmp/{}.jpg'.format(plot_csum), plot_img)
if opt.nas and opt.nas_stage > 0:
teacher_imgs = imgs.to('cuda:1')
with torch.no_grad():
inf_out, _ = teacher_model(teacher_imgs) # forward
# filter by obj confidence 0.05
teacher_pred = non_max_suppression_teacher(
inf_out, conf_thres=0.05, iou_thres=0.6, merge=False
) # (x1, y1, x2, y2, conf, cls) in resized image size
teacher_targets = teacher2targets(teacher_pred,
teacher_imgs)
# print('---> teacher_pred', teacher_pred)
# print('---> targets', targets)
# print('---> teacher_targets', teacher_targets)
# TODO: apply soft label loss
teacher_loss, teacher_loss_items = compute_teacher_loss(
pred, teacher_targets.to(device),
model) # loss scaled by batch_size
# print("===> origin loss", loss, loss_items)
# print("===> teacher loss", teacher_loss, teacher_loss_items)
teacher_loss_scale = 2.0
loss += teacher_loss * teacher_loss_scale
loss_items += teacher_loss_items * teacher_loss_scale
########## the targets and teacher predictions are matched, but they both can not be restored to the image, need TODO!! ###########
# assert len(teacher_pred) == imgs.size()[0]
# for i, (det, plot_img) in enumerate(zip(teacher_pred, imgs.detach().cpu().numpy())):
# plot_img = np.transpose(plot_img, (1,2,0))
# plot_img = np.uint8(plot_img * 255.0)
# plot_csum += 1
# cv2.imwrite('./tmp/{}.jpg'.format(plot_csum), plot_img)
# plot_img = cv2.imread('./tmp/{}.jpg'.format(plot_csum))
# for tgt in targets.detach().cpu().numpy():
# _, tgt_class_id, c_x, c_y, c_w, c_h = tgt
# c_x, c_y, c_w, c_h = float(c_x), float(c_y), float(c_w), float(c_h)
# c_x, c_y, c_w, c_h = c_x * plot_img.shape[1], c_y * plot_img.shape[0], c_w * plot_img.shape[1], c_h * plot_img.shape[0]
# cv2.rectangle(plot_img, (int(c_x - c_w / 2), int(c_y - c_h / 2)), (int(c_x + c_w / 2), int(c_y + c_h / 2)), (0,0,255), 2)
# print('===> ', int(c_x - c_w / 2), int(c_y - c_h / 2), int(c_x + c_w / 2), int(c_y + c_h / 2), tgt_class_id)
# if det is not None:
# det = det.detach().cpu().numpy()
# for each_b in det:
# pass
# cv2.rectangle(plot_img, (int(each_b[0]), int(each_b[1])), (int(each_b[2]), int(each_b[3])), (255,0,0), 2)
# print('---> ', int(each_b[0]), int(each_b[1]), int(each_b[2]), int(each_b[3]), float(each_b[4]), int(each_b[5]))
# cv2.imwrite('./tmp/{}.jpg'.format(plot_csum), plot_img)
# Backward
# scaler.scale(loss).backward()
loss.backward()
# Optimize
if ni % accumulate == 0:
# scaler.step(optimizer) # optimizer.step
# scaler.update()
optimizer.step()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if ni < 3:
f = str(log_dir / ('train_batch%g.jpg' % ni)) # filename
result = plot_images(images=imgs,
targets=targets,
paths=paths,
fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f,
result,
dataformats='HWC',
global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema:
ema.update_attr(
model,
include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
if opt.nas:
# only evaluate the super network
ema.ema.nas_stage = 0
results, maps, times = test.test(opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir)
if opt.nas:
ema.ema.nas_stage = opt.nas_stage
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results +
'\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' %
(results_file, opt.bucket, opt.name))
# Tensorboard
if tb_writer:
tags = [
'train/giou_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/giou_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
tb_writer.add_scalar(tag, x, epoch)
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'training_results':
f.read(),
'model':
ema.ema,
'optimizer':
None if final_epoch else optimizer.state_dict()
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
n = ('_'
if len(opt.name) and not opt.name.isnumeric() else '') + opt.name
fresults, flast, fbest = 'results%s.txt' % n, wdir / f'last{n}.pt', wdir / f'best{n}.pt'
for f1, f2 in zip([wdir / 'last.pt', wdir / 'best.pt', 'results.txt'],
[flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
if str(f2).endswith('.pt'): # is *.pt
strip_optimizer(f2) # strip optimizer
os.system(
'gsutil cp %s gs://%s/weights' %
(f2, opt.bucket)) if opt.bucket else None # upload
# Finish
if not opt.evolve:
plot_results(save_dir=log_dir) # save as results.png
logger.info('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
return results
def detect():
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
# Initialize
device = torch_utils.select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(
weights,
map_location=device)['model'].float().eval() # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Set Dataloader
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
raw_h, raw_w, _ = im0s.shape #
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = torch_utils.time_synchronized()
# Process detections
for i, det in enumerate(pred): # detections per image
p, s, im0 = path, '', im0s
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
submit_result.append({
'name':
path.split('/')[-1],
'category':
int(cls.item()) + 1,
'bbox': [
int(xyxy[0]),
int(xyxy[1]),
int(xyxy[2]),
int(xyxy[3])
],
'score':
conf.item()
})
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Save results (image with detections)
print('Done. (%.3fs)' % (time.time() - t0))
print(submit_result)
with open('resut_yolov5l.json', 'w') as fp:
json.dump(submit_result, fp, indent=4, ensure_ascii=False)
def detect(opt):
out, source, weights, imgsz = \
opt.output, opt.source, opt.weights, opt.img_size
parameters = opt.parameters
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# TRACKING_NUM = parameters["timelimit"] * 24 * 60
TRACKING_NUM = opt.record_length
save_img, view_img, save_txt = False, False, False
# get scenario
scenario = Path(opt.output).name
if scenario == 'beam':
TRACKING_NUM = 720 * 60
# Initialize
device = torch_utils.select_device(opt.device)
#print(device)
if opt.camera_id:
out += "/{}".format(opt.camera_id)
if os.path.exists(out):
shutil.rmtree(out, ignore_errors=True) # delete output folder
if not os.path.exists(out):
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(
weights,
map_location=device)['model'].float().eval() # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
# if scenario in {"damper"}:
# classify = True
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101_damper.pt',
map_location=device)) # load weights
modelc.to(device).eval()
status_classify = True if scenario in ["damper"] else False
if status_classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101_damper.pt',
map_location=device)) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
# fps = vid_cap.get(cv2.CAP_PROP_FPS)
temp_cap = cv2.VideoCapture(source)
w = int(temp_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(temp_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
temp_cap.release()
fps = 3
sp = out + "/" + time.strftime(
"%Y%m%d%H%M%S") + "_" + opt.online_save_name
print(sp)
stream_writer = cv2.VideoWriter(sp, fourcc, fps, (w, h))
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
Tracking = TRACKING_NUM
Start_tracking = False
for idx, data in enumerate(
dataset
): # path;fix size img;origin img;the video capture if there are video file
path, img, im0s, vid_cap = data
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float(
) # uint8 to fp16/32, half precise computation to accelerate speed;
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(
pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms) # pred: xyxy,prob,class
t2 = torch_utils.time_synchronized()
if scenario in ["rock"]:
pred = [x[x[:, 4] > 0.6] for x in pred if x is not None]
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
if status_classify:
pred = status_classifier(pred, modelc, img, im0s)
# Process detections
IS_VIOLATION = False
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
res, det, im0 = violation_detection(im0,
scenario=scenario,
detections=det,
names=names,
parameters=parameters)
# print(res)
IS_VIOLATION = res["violation"]
if scenario == "beam":
if res["工人"] or res["试探员"]:
IS_VIOLATION = res["violation"]
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
#im0 = im0[int(xyxy[1]):int(xyxy[3]), int(xyxy[0]):int(xyxy[2])]
# plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
# Print time (inference + NMS)
# print('%sDone. (%.3fs)' % (s, t2 - t1))
if det is None or len(det) == 0:
if scenario == "sensor":
IS_VIOLATION = True
res = {
"violation": 1,
"传感器位置错误": 0,
"无传感器": 0,
"传感器数目不足": 0,
"传感器离顶过近": 0,
"传感器离墙过近": 0,
"传感器悬挂过低": 0,
}
elif scenario == "Huge_rock":
IS_VIOLATION = False
res = {"violation": 0, "大型煤块": 0}
elif scenario == "beam":
IS_VIOLATION = False
res = {
"violation": 0,
"前探梁满足": 0,
"前探梁": 0,
"工人": 0,
"试探员": 0,
"液压机": 0
}
elif scenario == "damper":
IS_VIOLATION = True
res = {"violation": 1, "未检测到风门": 1, "风门打开": 0}
if IS_VIOLATION:
Start_tracking = True
# Print time (inference + NMS)
#print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
if Start_tracking:
if Tracking > 0:
Tracking -= 1
stream_writer.write(im0)
save_notes(sp + ".json", res)
else:
Tracking = TRACKING_NUM
Start_tracking = False
if isinstance(stream_writer, cv2.VideoWriter):
stream_writer.release()
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
# fps = vid_cap.get(cv2.CAP_PROP_FPS)
temp_cap = cv2.VideoCapture(source)
w = int(temp_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(temp_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
temp_cap.release()
fps = 6
sp = out + "/" + time.strftime(
"%Y%m%d%H%M%S") + "_" + opt.online_save_name
print(sp)
stream_writer = cv2.VideoWriter(
sp, fourcc, fps, (w, h))
else:
pass
if idx % 720 == 0:
cv2.imwrite(sp + "_{}_.png".format(idx), im0)
# print("Tracking : ", Tracking)
# print("idx : ", idx)
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
#fourcc = 'mp4v' # output video codec
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
#fps = vid_cap.get(cv2.CAP_PROP_FPS)
fps = 3
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(save_path, fourcc, fps,
(w, h))
if IS_VIOLATION:
vid_writer.write(im0)
save_notes(save_path + ".json", res)
if idx % 720 == 0:
cv2.imwrite(save_path + "_{}_.png".format(idx),
im0)
#IS_VIOLATION = False
#if evidence_vid_path != evidence_save_path: # new evidence piece
# evidence_vid_path = evidence_save_path
# if isinstance(evidence_vid_writer, cv2.VideoWriter):
# evidence_vid_writer.release()
# fourcc = 'mp4v' # output video codec
# #fps = vid_cap.get(cv2.CAP_PROP_FPS)
# fps = 24
# w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# evidence_vid_writer = cv2.VideoWriter(evidence_save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
#IS_VIOLATION:
# evidence_vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def train(hyp, opt, device, tb_writer=None, wandb=None):
logger.info(f'Hyperparameters {hyp}')
log_dir = Path(tb_writer.log_dir) if tb_writer else Path(
opt.logdir) / 'evolve' # logging directory
wdir = log_dir / 'weights' # weights dfirectory
os.makedirs(wdir, exist_ok=True)
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = str(log_dir / 'results.txt')
epochs, batch_size, total_batch_size, weights, rank = \
opt.epochs_init, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Save run settings
with open(log_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(log_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc, names = (1, ['item']) if opt.single_cls else (int(
data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (
len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get('anchors'):
ckpt['model'].yaml['anchors'] = round(
hyp['anchors']) # force autoanchor
model = Model(opt.cfg or ckpt['model'].yaml, ch=3,
nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [
] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict,
model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
'Transferred %g/%g items from %s' %
(len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
if opt.reg_lambda != 0:
# the regularization is based on Synaptic Intelligence as described in the
# paper. ewcData is a list of two elements (best parametes, importance)
# while synData is a dictionary with all the trajectory data needed by SI
model.ewcData, model.synData = create_syn_data(model)
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
optimizer = optim.SGD(pg0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': pg1,
'weight_decay': hyp['weight_decay']
}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' %
(len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp[
'lrf']) + hyp['lrf'] # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Logging
if wandb and wandb.run is None:
id = ckpt.get('wandb_id') if 'ckpt' in locals() else None
wandb_run = wandb.init(config=opt,
resume="allow",
project="YOLOv5",
name=os.path.basename(log_dir),
id=id)
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (
weights, epochs)
shutil.copytree(wdir, wdir.parent /
f'weights_backup_epoch{start_epoch - 1}'
) # save previous weights
if epochs < start_epoch:
logger.info(
'%s has been trained for %g epochs. Fine-tuning for %g additional epochs.'
% (weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
all_test_dataloader = create_dataloader(test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images
and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers,
n_batch=-1)[0]
root = 'G:/projects/core50_350_1f/batches/'
paths = os.listdir(root)
train_paths = []
valid_paths = []
for p in paths:
if 'train' in p:
train_paths.append(root + p)
elif 'val' in p:
valid_paths.append(root + p)
else:
print(p)
# external_memory = ext_memory()
extMem = externalMemory()
for core_batch in range(11):
# Trainloader
if opt.reg_lambda != 0:
init_batch(model, model.ewcData, model.synData)
print(f'------------CORE50 itertaion №:{core_batch}------------')
external_files_path = extMem.file
if core_batch > 0:
train_path = [train_paths[core_batch], external_files_path]
else:
train_path = train_paths[core_batch]
extMem.update_memory(train_paths[core_batch],
update_iters=10 if core_batch == 0 else 1)
dataloader, dataset = create_dataloader(
train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers,
)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (
mlc, nc, opt.data, nc - 1)
testloader = create_dataloader(valid_paths[core_batch],
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
augment=False,
cache=opt.cache_images
and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers)[0] # testloader
#if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.IntTensor(labels[:, 0]) # classes
plot_labels(labels, save_dir=log_dir)
print(torch.bincount(c))
if tb_writer:
# tb_writer.add_hparams(hyp, {}) # causes duplicate https://github.com/ultralytics/yolov5/pull/384
tb_writer.add_histogram('classes', c, core_batch)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
# model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
model.names = names
# Start training
t0 = time.time()
nw = max(
round(hyp['warmup_epochs'] * nb),
1e3) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
logger.info(
'Image sizes %g train, %g test\n'
'Using %g dataloader workers\nLogging results to %s\n'
'Starting training for %g epochs...' %
(imgsz, imgsz_test, dataloader.num_workers, log_dir, epochs))
# update number of epochs to iterative training
if core_batch != 0:
epochs = opt.epochs_iter
# x_train, y_train = dataset.get_all_data()
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
mloss = torch.zeros(4, device=device) # mean losses
logger.info(
('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls',
'total', 'targets', 'img_size'))
# x_train_splitted = torch.split(x_train, 4)
# y_train_splitted = torch.split(y_train, 4)
# pbar = enumerate(zip(x_train_splitted, y_train_splitted))
pbar = enumerate(dataloader)
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, _, _
) in pbar: # batch -------------------------------------------------------------
# imgs = x_train[i * batch_size:(i + 1) * batch_size]
# targets = y_train[i * batch_size:(i + 1) * batch_size]
#
# # preprocess tensor to proper form
# # img, label = zip(imgs, targets) # transposed
# for i, l in enumerate(targets):
# l[:, 0] = i # add target image index for build_targets()
#
# imgs = torch.stack(imgs)
# targets = torch.cat(targets)
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
if opt.reg_lambda != 0:
pre_update(model, model.synData)
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi,
[hyp['warmup_momentum'], hyp['momentum']])
# # Multi-scale
# if opt.multi_scale:
# sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
# sf = sz / max(imgs.shape[2:]) # scale factor
# if sf != 1:
# ns = [math.ceil(x * sf / gs) * gs for x in
# imgs.shape[2:]] # new shape (stretched to gs-multiple)
# imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device),
model) # loss scaled by batch_size
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
if opt.reg_lambda != 0:
post_update(model, model.synData)
scaler.update()
optimizer.zero_grad()
# if ema:
# ema.update(model)
# Print
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9
if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 +
'%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem,
*mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# mAP
results, maps, times = test.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
# model=ema.ema,
model=model,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=log_dir,
plots=epoch == 0, # plot first and last
log_imgs=opt.log_imgs)
# wandb.log({'per class/AP per class': maps})
# Write
with open(results_file, 'a') as f:
f.write(
s + '%10.4g' * 7 % results +
'\n') # P, R, [email protected], [email protected], val_loss(box, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' %
(results_file, opt.bucket, opt.name))
# Log
tags = [
'train/giou_loss',
'train/obj_loss',
'train/cls_loss', # train loss
'metrics/precision',
'metrics/recall',
'metrics/mAP_0.5',
'metrics/mAP_0.5:0.95',
'val/giou_loss',
'val/obj_loss',
'val/cls_loss', # val loss
'x/lr0',
'x/lr1',
'x/lr2'
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if tb_writer:
tb_writer.add_scalar(tag, x, epoch) # tensorboard
if wandb:
wandb.log({tag: x}) # W&B
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
# Save model
save = not opt.nosave
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'training_results': f.read(),
# 'model': ema.ema,
'model': model,
'optimizer': optimizer.state_dict(),
'wandb_id': wandb_run.id if wandb else None
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
#consolidate_weights(model, cur_class)
if opt.reg_lambda != 0:
update_ewc_data(model, model.ewcData, model.synData, 0.001, 1)
if rank in [-1, 0]:
# Strip optimizers
n = opt.name if opt.name.isnumeric() else ''
fresults, flast, fbest = log_dir / f'results{n}.txt', wdir / f'last{n}.pt', wdir / f'best{n}.pt'
for f1, f2 in zip(
[wdir / 'last.pt', wdir / 'best.pt', results_file],
[flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
if str(f2).endswith('.pt'): # is *.pt
strip_optimizer(f2) # strip optimizer
os.system(
'gsutil cp %s gs://%s/weights' %
(f2, opt.bucket)) if opt.bucket else None # upload
# Finish
plot_results(save_dir=log_dir) # save as results.png
logger.info('%g epochs completed in %.3f hours.\n' %
(epoch - start_epoch + 1, (time.time() - t0) / 3600))
dist.destroy_process_group() if rank not in [-1, 0] else None
torch.cuda.empty_cache()
results, maps, times = test.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
#model=ema.ema,
model=model,
single_cls=opt.single_cls,
dataloader=all_test_dataloader,
save_dir=log_dir,
#plots=epoch == 0 or final_epoch, # plot first and last
log_imgs=opt.log_imgs,
verbose=True)
#wandb.log({'per class/AP per class All': maps[0]})
#tb_writer.add_scalar('per class/AP per class All', maps[0])
# Log
tags = [ # train loss
'test/precision', 'test/recall', 'test/mAP_0.5',
'test/mAP_0.5:0.95', 'test/giou_loss', 'test/obj_loss',
'test/cls_loss'
] # params
for x, tag in zip(list(results), tags):
if tb_writer:
tb_writer.add_scalar(tag, x, core_batch) # tensorboard
if wandb:
wandb.log({tag: x}) # W&B
return results
def test(data,
weights=None,
batch_size=16,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
merge=False):
# Initialize/load model and set device
if model is None:
training = False
device = torch_utils.select_device(opt.device, batch_size=batch_size)
half = device.type != 'cpu' # half precision only supported on CUDA
# Remove previous
for f in glob.glob('test_batch*.jpg'):
os.remove(f)
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights, map_location=device)['model'].float() # load to FP32
torch_utils.model_info(model)
model.fuse()
model.to(device)
if half:
model.half() # to FP16
# Multi-GPU disabled, incompatible with .half()
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
else: # called by train.py
training = True
device = next(model.parameters()).device # get model device
# half disabled https://github.com/ultralytics/yolov5/issues/99
half = False # device.type != 'cpu' and torch.cuda.device_count() == 1
if half:
model.half() # to FP16
# Configure
model.eval()
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for [email protected]:0.95
# iouv = iouv[0].view(1) # comment for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if dataloader is None: # not training
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
merge = opt.merge # use Merge NMS
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
dataset = LoadImagesAndLabels(path,
imgsz,
batch_size,
rect=True, # rectangular inference
single_cls=opt.single_cls, # single class mode
pad=0.5) # padding
batch_size = min(batch_size, len(dataset))
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
dataloader = DataLoader(dataset,
batch_size=batch_size,
num_workers=nw,
pin_memory=True,
collate_fn=dataset.collate_fn)
seen = 0
names = model.names if hasattr(model, 'names') else model.module.names
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', '[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
whwh = torch.Tensor([width, height, width, height]).to(device)
# Disable gradients
with torch.no_grad():
# Run model
t = torch_utils.time_synchronized()
inf_out, train_out = model(img, augment=augment) # inference and training outputs
t0 += torch_utils.time_synchronized() - t
# Compute loss
if training: # if model has loss hyperparameters
loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls
# Run NMS
t = torch_utils.time_synchronized()
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, merge=merge)
t1 += torch_utils.time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
seen += 1
if pred is None:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
continue
# Append to text file
# with open('test.txt', 'a') as file:
# [file.write('%11.5g' * 7 % tuple(x) + '\n') for x in pred]
# Clip boxes to image bounds
clip_coords(pred, (height, width))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(Path(paths[si]).stem.split('_')[-1])
box = pred[:, :4].clone() # xyxy
scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape
box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({'image_id': image_id,
'category_id': coco91class[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5]) * whwh
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero().view(-1) # prediction indices
pi = (cls == pred[:, 5]).nonzero().view(-1) # target indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices
# Append detections
for j in (ious > iouv[0]).nonzero():
d = ti[i[j]] # detected target
if d not in detected:
detected.append(d)
correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(detected) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# Plot images
if batch_i < 1:
f = 'test_batch%g_gt.jpg' % batch_i # filename
plot_images(img, targets, paths, f, names) # ground truth
f = 'test_batch%g_pred.jpg' % batch_i
plot_images(img, output_to_target(output, width, height), paths, f, names) # predictions
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats):
p, r, ap, f1, ap_class = ap_per_class(*stats)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%12.3g' * 6 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if verbose and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
# Save JSON
if save_json and map50 and len(jdict):
imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
f = 'detections_val2017_%s_results.json' % \
(weights.split(os.sep)[-1].replace('.pt', '') if weights else '') # filename
print('\nCOCO mAP with pycocotools... saving %s...' % f)
with open(f, 'w') as file:
json.dump(jdict, file)
try:
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
# https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api
cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
cocoEval.params.imgIds = imgIds # image IDs to evaluate
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
map, map50 = cocoEval.stats[:2] # update results ([email protected]:0.95, [email protected])
except:
print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. '
'See https://github.com/cocodataset/cocoapi/issues/356')
# Return results
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
def detect_json(options, save_img=False):
"""
Returns a json file which includes the results of the inference:
Form of the json
Video: {0: [{},{}]}:
Image:
"""
opt = options
out, source, weights, view_img, save_txt, imgsz = \
opt["output"], opt["source"], opt["weights"], opt["view_img"], opt["save_txt"], opt["img_size"]
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(opt["device"])
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.names if hasattr(model, 'names') else model.modules.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
colors[0] = (255, 0, 0)
colors[1] = (0, 255, 0)
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
result_dic = dict()
img_counter = 0
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt["augment"])[0]
# Apply NMS
pred = non_max_suppression(pred,
opt["conf_thres"],
opt["iou_thres"],
classes=opt["classes"],
agnostic=opt["agnostic_nms"])
t2 = torch_utils.time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
#No we create the result dict
result_list_one_frame = []
for *xyxy, conf, cls in det:
single_inference_dic = {}
if int(cls) == 0:
label = "mask"
elif int(cls) == 1:
label = "no mask"
else:
sys.exit("Label not known")
single_inference_dic["xywh"] = (
xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
single_inference_dic["xywh"] = [
float(i) for i in single_inference_dic["xywh"]
]
single_inference_dic["xyxy"] = ((float(xyxy[0]),
float(xyxy[1])),
(float(xyxy[2]),
float(xyxy[3])))
single_inference_dic["map"] = float(conf)
single_inference_dic["label"] = str(label)
result_list_one_frame.append(single_inference_dic)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
img_counter += 1
result_dic[img_counter] = result_list_one_frame
print(result_dic)
print('Done. (%.3fs)' % (time.time() - t0))
return result_dic
def train(hyp):
epochs = opt.epochs # 300
batch_size = opt.batch_size # 64
weights = opt.weights # initial training weights
# Configure
init_seeds(1)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
train_path = data_dict['train']
test_path = data_dict['val']
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
# Remove previous results
for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):
os.remove(f)
# Create model
model = Model(opt.cfg, nc=data_dict['nc']).to(device)
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_parameters():
if v.requires_grad:
if '.bias' in k:
pg2.append(v) # biases
elif '.weight' in k and '.bn' not in k:
pg1.append(v) # apply weight decay
else:
pg0.append(v) # all else
optimizer = optim.Adam(pg0, lr=hyp['lr0']) if opt.adam else \
optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
lf = lambda x: (((1 + math.cos(x * math.pi / epochs)) / 2) ** 1.0) * 0.9 + 0.1 # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Load Model
google_utils.attempt_download(weights)
start_epoch, best_fitness = 0, 0.0
if weights.endswith('.pt'): # pytorch format
ckpt = torch.load(weights, map_location=device) # load checkpoint
# load model
try:
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items()
if model.state_dict()[k].shape == v.shape} # to FP32, filter
model.load_state_dict(ckpt['model'], strict=False)
except KeyError as e:
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \
"Please delete or update %s and try again, or use --weights '' to train from scratch." \
% (opt.weights, opt.cfg, opt.weights, opt.weights)
raise KeyError(s) from e
# load optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# load results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# epochs
start_epoch = ckpt['epoch'] + 1
if epochs < start_epoch:
print('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
(opt.weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt
# Mixed precision training https://github.com/NVIDIA/apex
if mixed_precision:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)
scheduler.last_epoch = start_epoch - 1 # do not move
# https://discuss.pytorch.org/t/a-problem-occured-when-resuming-an-optimizer/28822
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Initialize distributed training
if device.type != 'cpu' and torch.cuda.device_count() > 1 and torch.distributed.is_available():
dist.init_process_group(backend='nccl', # distributed backend
init_method='tcp://127.0.0.1:9999', # init method
world_size=1, # number of nodes
rank=0) # node rank
model = torch.nn.parallel.DistributedDataParallel(model)
# pip install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
# Trainloader
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Correct your labels or your model.' % (mlc, nc, opt.cfg)
# Testloader
testloader = create_dataloader(test_path, imgsz_test, batch_size, gs, opt,
hyp=hyp, augment=False, cache=opt.cache_images, rect=True)[0]
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
model.names = data_dict['names']
# Class frequency
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1.
# model._initialize_biases(cf.to(device))
if tb_writer:
plot_labels(labels)
tb_writer.add_histogram('classes', c, 0)
# Check anchors
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
# Exponential moving average
ema = torch_utils.ModelEMA(model)
# Start training
t0 = time.time()
nb = len(dataloader) # number of batches
n_burn = max(3 * nb, 1e3) # burn-in iterations, max(3 epochs, 1k iterations)
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0) # 'P', 'R', 'mAP', 'F1', 'val GIoU', 'val Objectness', 'val Classification'
print('Image sizes %g train, %g test' % (imgsz, imgsz_test))
print('Using %g dataloader workers' % dataloader.num_workers)
print('Starting training for %g epochs...' % epochs)
# torch.autograd.set_detect_anomaly(True)
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if dataset.image_weights:
w = model.class_weights.cpu().numpy() * (1 - maps) ** 2 # class weights
image_weights = labels_to_image_weights(dataset.labels, nc=nc, class_weights=w)
dataset.indices = random.choices(range(dataset.n), weights=image_weights, k=dataset.n) # rand weighted idx
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
print(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'GIoU', 'obj', 'cls', 'total', 'targets', 'img_size'))
pbar = tqdm(enumerate(dataloader), total=nb) # progress bar
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device).float() / 255.0 # uint8 to float32, 0 - 255 to 0.0 - 1.0
# Burn-in
if ni <= n_burn:
xi = [0, n_burn] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [0.1 if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
pred = model(imgs)
# Loss
loss, loss_items = compute_loss(pred, targets.to(device), model)
if not torch.isfinite(loss):
print('WARNING: non-finite loss, ending training ', loss_items)
return results
# Backward
if mixed_precision:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# Optimize
if ni % accumulate == 0:
optimizer.step()
optimizer.zero_grad()
ema.update(model)
# Print
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 + '%10.4g' * 6) % (
'%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if ni < 3:
f = 'train_batch%g.jpg' % ni # filename
result = plot_images(images=imgs, targets=targets, paths=paths, fname=f)
if tb_writer and result is not None:
tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
scheduler.step()
# mAP
ema.update_attr(model)
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(opt.data,
batch_size=batch_size,
imgsz=imgsz_test,
save_json=final_epoch and opt.data.endswith(os.sep + 'coco.yaml'),
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader)
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results + '\n') # P, R, mAP, F1, test_losses=(GIoU, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp results.txt gs://%s/results/results%s.txt' % (opt.bucket, opt.name))
# Tensorboard
if tb_writer:
tags = ['train/giou_loss', 'train/obj_loss', 'train/cls_loss',
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/F1',
'val/giou_loss', 'val/obj_loss', 'val/cls_loss']
for x, tag in zip(list(mloss[:-1]) + list(results), tags):
tb_writer.add_scalar(tag, x, epoch)
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # fitness_i = weighted combination of [P, R, mAP, F1]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'training_results': f.read(),
'model': ema.ema,
'optimizer': None if final_epoch else optimizer.state_dict()}
# Save last, best and delete
torch.save(ckpt, last)
if (best_fitness == fi) and not final_epoch:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
# Strip optimizers
n = ('_' if len(opt.name) and not opt.name.isnumeric() else '') + opt.name
fresults, flast, fbest = 'results%s.txt' % n, wdir + 'last%s.pt' % n, wdir + 'best%s.pt' % n
for f1, f2 in zip([wdir + 'last.pt', wdir + 'best.pt', 'results.txt'], [flast, fbest, fresults]):
if os.path.exists(f1):
os.rename(f1, f2) # rename
ispt = f2.endswith('.pt') # is *.pt
strip_optimizer(f2) if ispt else None # strip optimizer
os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket and ispt else None # upload
# Finish
if not opt.evolve:
plot_results() # save as results.png
print('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
dist.destroy_process_group() if device.type != 'cpu' and torch.cuda.device_count() > 1 else None
torch.cuda.empty_cache()
return results
def detect_demo(options, save_img=False):
opt = options
out, source, weights, view_img, save_txt, imgsz = \
opt["output"], opt["source"], opt["weights"], opt["view_img"], opt["save_txt"], opt["img_size"]
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(opt["device"])
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.names if hasattr(model, 'names') else model.modules.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
keep_overlay = False
overlay_to_use = None
overlay_to_use_old = None
iteration_subcounter = 0
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt["augment"])[0]
# Apply NMS
pred = non_max_suppression(pred,
opt["conf_thres"],
opt["iou_thres"],
classes=opt["classes"],
agnostic=opt["agnostic_nms"])
t2 = torch_utils.time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
#Initial load of the diffent overlays
path_to_ok = "demo_files/ok_sign_final.png"
overlay_ok = cv2.imread(path_to_ok)
overlay_ok = cv2.resize(overlay_ok, (1920, 1080))
path_to_stop = "demo_files/stop_sign_final.png"
overlay = cv2.imread(path_to_stop)
trans_mask = overlay[:, :, 2] == 0
overlay[trans_mask] = [255, 255, 255]
overlay_not_ok = cv2.cvtColor(overlay, cv2.COLOR_BGRA2BGR)
overlay_not_ok = cv2.resize(overlay_not_ok, (1920, 1080))
path_to_attention = "demo_files/attention_sign_final.png"
overlay_attention = cv2.imread(path_to_attention)
overlay_attention = cv2.resize(overlay_attention, (1920, 1080))
# Process detections
for i, det in enumerate(pred): # detections per image
print("tetes")
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
no_mask_true = False
mask_true = False
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(save_path[:save_path.rfind('.')] + '.txt',
'a') as file:
file.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if int(cls) == 0:
mask_true = True
elif int(cls) == 1:
no_mask_true = True
else:
sys.exit("Label not known")
if save_img or view_img: # Add bbox to image
label = None
if int(cls) == 0:
label = "mask"
color = (0, 255, 0)
elif int(cls) == 1:
label = "No mask"
color = (0, 0, 255)
else:
sys.exit("Label not known")
plot_one_box(xyxy,
im0,
label=label,
color=color,
line_thickness=3)
print(keep_overlay, iteration_subcounter, overlay_to_use,
overlay_to_use_old)
frame_to_wait = 100
if (keep_overlay
== True) and (iteration_subcounter < frame_to_wait):
overlay_to_use = overlay_to_use_old
iteration_subcounter += 1
elif (keep_overlay
== True) and (iteration_subcounter >= frame_to_wait):
iteration_subcounter = 0
keep_overlay = False
if (mask_true == True) and (no_mask_true
== False) and (keep_overlay == False):
overlay_to_use = "ok"
elif (no_mask_true == True) and (keep_overlay == False):
overlay_to_use = "not_ok"
elif (mask_true == False) and (no_mask_true
== False) and (keep_overlay
== False):
overlay_to_use = "attention"
else:
pass
#
if overlay_to_use != overlay_to_use_old:
keep_overlay = True
iteration_subcounter = 0
overlay_to_use_old = overlay_to_use
else:
overlay_to_use_old = overlay_to_use
if overlay_to_use == "ok":
overlay = overlay_ok
elif overlay_to_use == "not_ok":
overlay = overlay_not_ok
elif overlay_to_use == "attention":
overlay = overlay_attention
else:
print(f"Overlay not known {overlay_to_use}")
#Including the current images
im0 = cv2.resize(im0, (1920, 1080))
#include_small_image = True
#info_screen_small = False
if options["info_screen_small"] == False:
s_img = im0.copy()
s_img = cv2.resize(s_img, (1080, 1920))
print(s_img.shape, overlay.shape)
s_img = cv2.resize(s_img, (480, 360))
print(s_img.shape, overlay.shape)
x_offset = 10
y_offset = 10
if options["include_small_image"] == True:
overlay[y_offset:y_offset + s_img.shape[0],
x_offset:x_offset + s_img.shape[1]] = s_img
im0 = overlay
else:
s_img = overlay.copy()
s_img = cv2.resize(s_img, (1080, 1920))
print(s_img.shape, overlay.shape)
s_img = cv2.resize(s_img, (480, 360))
print(s_img.shape, im0.shape)
x_offset = 10
y_offset = 10
if options["include_small_image"] == True:
im0[y_offset:y_offset + s_img.shape[0],
x_offset:x_offset + s_img.shape[1]] = s_img
im0 = im0
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.namedWindow("window", cv2.WND_PROP_FULLSCREEN)
cv2.setWindowProperty("window", cv2.WND_PROP_FULLSCREEN,
cv2.WINDOW_FULLSCREEN)
#im0 = cv2.resize(im0,(1920,1080))
cv2.imshow("window", im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
save_img = True
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt["fourcc"]),
fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % save_path)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
from utils import google_utils
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
print(opt)
# Input
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection
# Load PyTorch model
google_utils.attempt_download(opt.weights)
model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
model.eval()
model.model[-1].export = True # set Detect() layer export=True
_ = model(img) # dry run
# TorchScript export
try:
f = opt.weights.replace('.pt', '.torchscript') # filename
ts = torch.jit.trace(model, img)
ts.save(f)
print('TorchScript export success, saved as %s' % f)
except Exception as e:
print('TorchScript export failed: %s' % e)
# ONNX export
def detect(save_img=True, view_img=False):
logging.info('##################### start detecting #####################')
try:
out, source, weights, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith(
'rtsp') or source.startswith('http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(opt.device)
if not os.path.exists(out):
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(
weights, map_location=device)['model'].float() # load to FP32
model.to(device).eval()
if half:
model.half() # to FP16
# Set Dataloader
if webcam:
# view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
init_x = [None] * len(names)
end_x = [None] * len(names)
for path, img, im0s, vid_cap in dataset:
torch.cuda.empty_cache()
output = out + '/' + time.strftime('%Y-%m-%d', time.localtime())
if not os.path.exists(output):
os.makedirs(output)
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = torch_utils.time_synchronized()
print('pred:', pred)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
logging.info(
'-------------------------------------------------\n%s'
% det)
now = time.strftime('%Y-%m-%d %H.%M.%S', time.localtime())
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
c = int(c.cpu().detach().numpy())
if init_x[c] is None:
init_x[c] = end_x[c] = det[det[:,
-1] == c][:,
0].min()
else:
# Write results
end_x[c] = det[det[:, -1] == c][:, 0].min()
if end_x[c] <= init_x[c] and end_x[c] < 1100:
retval, buffer = cv2.imencode('.jpeg', im0)
pic_str = base64.b64encode(buffer).decode()
status, timestamp, result = is_open(
cls, pic_str)
if result is not None:
logging.info(
'result: %s' %
json.dumps(result, ensure_ascii=False))
if status:
logging.info('************ 开门 ***********')
res_data = send_command(
server_url,
'open',
number=result['words_result']['number']
if c == 2 else None)
logging.info(
'status code: %s' % json.dumps(
res_data, ensure_ascii=False))
init_x[c] = end_x[c] = None
time.sleep(10)
elif end_x[c] > init_x[c]:
if init_x[c] < 700 and end_x[c] - init_x[
c] > 300:
logging.info(
'************ 关门 ************')
res_data = send_command(
server_url, 'close')
logging.info(
'status code: %s' % json.dumps(
res_data, ensure_ascii=False))
logging.info(
'init_x: %.3f, end_x: %.3f' %
(init_x[c].cpu().detach().numpy(),
end_x[c].cpu().detach().numpy()))
init_x[c] = end_x[c] = None
time.sleep(10)
for *xyxy, conf, cls in det[det[:, -1] == c]:
cls = int(cls.cpu().detach().numpy())
# # xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
# if xyxy[0] <= init_x[cls] and xyxy[0] < 1100:
# end_x[cls] = xyxy[0]
# retval, buffer = cv2.imencode('.jpeg', im0)
# pic_str = base64.b64encode(buffer).decode()
# status, timestamp, result = is_open(cls, pic_str)
# if result is not None:
# logging.info('result: %s' % json.dumps(result, ensure_ascii=False))
# if status:
# logging.info('************ 开门 ***********')
# res_data = send_command(
# server_url,
# 'open',
# number=result['words_result']['number'] if cls == 2 else None)
# logging.info('status code: %s' % json.dumps(res_data, ensure_ascii=False))
# init_x[cls] = end_x[cls] = None
# time.sleep(10)
# elif xyxy[0] > init_x[cls]:
# end_x[cls] = xyxy[0]
# if init_x[cls] < 700 and end_x[cls] - init_x[cls] > 300:
# logging.info('************ 关门 ************')
# res_data = send_command(server_url, 'close')
# logging.info('status code: %s' % json.dumps(res_data, ensure_ascii=False))
# logging.info(
# 'init_x: %.3f, end_x: %.3f' % (init_x[cls].cpu().detach().numpy(), end_x[cls].cpu().detach().numpy()))
# init_x[cls] = end_x[cls] = None
# time.sleep(10)
# xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
# with open(txt_path + '.txt', 'a') as f:
# f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[cls], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
if save_img:
save_path = os.path.join(output, now + '.jpg')
logging.info('保存图片至：%s' % save_path)
cv2.imencode('.jpg', im0)[1].tofile(save_path)
else:
for cls in opt.classes:
if end_x[cls] is not None and init_x[
cls] is not None and end_x[cls] > init_x[
cls] and init_x[cls] < 700 and (
end_x[cls] - init_x[cls] > 300
or end_x[cls] > 1200):
logging.info(
'*********** 首次从有物体转换为无物体，识别关门 ***********')
res_data = send_command(server_url, 'close')
logging.info(
'status code: %s' %
json.dumps(res_data, ensure_ascii=False))
logging.info('init_x: %.3f, end_x: %.3f' %
(init_x[cls].cpu().detach().numpy(),
end_x[cls].cpu().detach().numpy()))
time.sleep(5)
else:
time.sleep(0.5)
init_x[cls] = end_x[cls] = None
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
# cv2.namedWindow("result", cv2.WINDOW_NORMAL)
cv2.imshow(p, im0)
if cv2.waitKey(300) & 0xFF == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
# if save_img:
# if dataset.mode == 'images':
# cv2.imwrite(save_path, im0)
# else:
# if vid_path != save_path: # new video
# vid_path = save_path
# if isinstance(vid_writer, cv2.VideoWriter):
# vid_writer.release() # release previous video writer
#
# fps = vid_cap.get(cv2.CAP_PROP_FPS)
# w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
# vid_writer.write(im0)
# if save_txt or save_img:
# print('Results saved to %s' % os.getcwd() + os.sep + out)
# if platform == 'darwin': # MacOS
# os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
detect()
except RuntimeError:
logging.exception('detect 报错')
return
except Exception:
# traceback.print_exc()
logging.exception('detect 报错')
torch.cuda.empty_cache()
time.sleep(5)
detect()
def train(hyp, opt, device, tb_writer=None, wandb=None):
logger.info(f"Hyperparameters {hyp}")
save_dir, epochs, batch_size, total_batch_size, weights, rank = (
Path(opt.save_dir),
opt.epochs,
opt.batch_size,
opt.total_batch_size,
opt.weights,
opt.global_rank,
)
# Directories
wdir = save_dir / "weights"
wdir.mkdir(parents=True, exist_ok=True) # make dir
last = wdir / "last.pt"
best = wdir / "best.pt"
results_file = save_dir / "results.txt"
# Save run settings
with open(save_dir / "hyp.yaml", "w") as f:
yaml.dump(hyp, f, sort_keys=False)
with open(save_dir / "opt.yaml", "w") as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
plots = not opt.evolve # create plots
cuda = device.type != "cpu"
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict["train"]
test_path = data_dict["val"]
nc, names = (
(1, ["item"]) if opt.single_cls else (int(data_dict["nc"]), data_dict["names"])
) # number classes, names
assert len(names) == nc, "%g names found for nc=%g dataset in %s" % (
len(names),
nc,
opt.data,
) # check
# Model
pretrained = weights.endswith(".pt")
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get("anchors"):
ckpt["model"].yaml["anchors"] = round(hyp["anchors"]) # force autoanchor
model = Model(opt.cfg or ckpt["model"].yaml, ch=3, nc=nc).to(device) # create
exclude = ["anchor"] if opt.cfg or hyp.get("anchors") else [] # exclude keys
state_dict = ckpt["model"].float().state_dict() # to FP32
state_dict = intersect_dicts(
state_dict, model.state_dict(), exclude=exclude
) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info(
"Transferred %g/%g items from %s"
% (len(state_dict), len(model.state_dict()), weights)
) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print("freezing %s" % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(
round(nbs / total_batch_size), 1
) # accumulate loss before optimizing
hyp["weight_decay"] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
if opt.adam:
optimizer = optim.Adam(
pg0, lr=hyp["lr0"], betas=(hyp["momentum"], 0.999)
) # adjust beta1 to momentum
else:
optimizer = optim.SGD(
pg0, lr=hyp["lr0"], momentum=hyp["momentum"], nesterov=True
)
optimizer.add_param_group(
{"params": pg1, "weight_decay": hyp["weight_decay"]}
) # add pg1 with weight_decay
optimizer.add_param_group({"params": pg2}) # add pg2 (biases)
logger.info(
"Optimizer groups: %g .bias, %g conv.weight, %g other"
% (len(pg2), len(pg1), len(pg0))
)
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = (
lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp["lrf"])
+ hyp["lrf"]
) # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Logging
if wandb and wandb.run is None:
opt.hyp = hyp # add hyperparameters
wandb_run = wandb.init(
config=opt,
resume="allow",
project="YOLOv5" if opt.project == "runs/train" else Path(opt.project).stem,
name=save_dir.stem,
id=ckpt.get("wandb_id") if "ckpt" in locals() else None,
)
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt["optimizer"] is not None:
optimizer.load_state_dict(ckpt["optimizer"])
best_fitness = ckpt["best_fitness"]
# Results
if ckpt.get("training_results") is not None:
with open(results_file, "w") as file:
file.write(ckpt["training_results"]) # write results.txt
# Epochs
start_epoch = ckpt["epoch"] + 1
if opt.resume:
assert (
start_epoch > 0
), "%s training to %g epochs is finished, nothing to resume." % (
weights,
epochs,
)
if epochs < start_epoch:
logger.info(
"%s has been trained for %g epochs. Fine-tuning for %g additional epochs."
% (weights, ckpt["epoch"], epochs)
)
epochs += ckpt["epoch"] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [
check_img_size(x, gs) for x in opt.img_size
] # verify imgsz are gs-multiples
# DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info("Using SyncBatchNorm()")
# EMA
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)
# Trainloader
dataloader, dataset = create_dataloader(
train_path,
imgsz,
batch_size,
gs,
opt,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=rank,
world_size=opt.world_size,
workers=opt.workers,
)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert (
mlc < nc
), "Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g" % (
mlc,
nc,
opt.data,
nc - 1,
)
# Process 0
if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates
testloader = create_dataloader(
test_path,
imgsz_test,
total_batch_size,
gs,
opt,
hyp=hyp,
cache=opt.cache_images and not opt.notest,
rect=True,
rank=-1,
world_size=opt.world_size,
workers=opt.workers,
)[
0
] # testloader
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
if plots:
plot_labels(labels, save_dir=save_dir)
if tb_writer:
tb_writer.add_histogram("classes", c, 0)
if wandb:
wandb.log(
{
"Labels": [
wandb.Image(str(x), caption=x.name)
for x in save_dir.glob("*labels*.png")
]
}
)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz)
# Model parameters
hyp["cls"] *= nc / 80.0 # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(
device
) # attach class weights
model.names = names
# Start training
t0 = time.time()
nw = max(
round(hyp["warmup_epochs"] * nb), 1000
) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
logger.info(
"Image sizes %g train, %g test\n"
"Using %g dataloader workers\nLogging results to %s\n"
"Starting training for %g epochs..."
% (imgsz, imgsz_test, dataloader.num_workers, save_dir, epochs)
)
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if rank in [-1, 0]:
cw = (
model.class_weights.cpu().numpy() * (1 - maps) ** 2
) # class weights
iw = labels_to_image_weights(
dataset.labels, nc=nc, class_weights=cw
) # image weights
dataset.indices = random.choices(
range(dataset.n), weights=iw, k=dataset.n
) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = (
torch.tensor(dataset.indices)
if rank == 0
else torch.zeros(dataset.n)
).int()
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(
("\n" + "%10s" * 8)
% ("Epoch", "gpu_mem", "box", "obj", "cls", "total", "targets", "img_size")
)
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs,
targets,
paths,
_,
) in (
pbar
): # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = (
imgs.to(device, non_blocking=True).float() / 255.0
) # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1, np.interp(ni, xi, [1, nbs / total_batch_size]).round()
)
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x["lr"] = np.interp(
ni,
xi,
[
hyp["warmup_bias_lr"] if j == 2 else 0.0,
x["initial_lr"] * lf(epoch),
],
)
if "momentum" in x:
x["momentum"] = np.interp(
ni, xi, [hyp["warmup_momentum"], hyp["momentum"]]
)
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [
math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = F.interpolate(
imgs, size=ns, mode="bilinear", align_corners=False
)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device), model
) # loss scaled by batch_size
if rank != -1:
loss *= (
opt.world_size
) # gradient averaged between devices in DDP mode
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = "%.3gG" % (
torch.cuda.memory_reserved() / 1e9
if torch.cuda.is_available()
else 0
) # (GB)
s = ("%10s" * 2 + "%10.4g" * 6) % (
"%g/%g" % (epoch, epochs - 1),
mem,
*mloss,
targets.shape[0],
imgs.shape[-1],
)
pbar.set_description(s)
# Plot
if plots and ni < 3:
f = save_dir / f"train_batch{ni}.jpg" # filename
plot_images(images=imgs, targets=targets, paths=paths, fname=f)
# if tb_writer:
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
elif plots and ni == 3 and wandb:
wandb.log(
{
"Mosaics": [
wandb.Image(str(x), caption=x.name)
for x in save_dir.glob("train*.jpg")
]
}
)
# end batch ------------------------------------------------------------------------------------------------
# end epoch ----------------------------------------------------------------------------------------------------
# Scheduler
lr = [x["lr"] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema:
ema.update_attr(
model, include=["yaml", "nc", "hyp", "gr", "names", "stride"]
)
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(
opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
plots=plots and final_epoch,
log_imgs=opt.log_imgs if wandb else 0,
)
# Write
with open(results_file, "a") as f:
f.write(
s + "%10.4g" * 7 % results + "\n"
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
if len(opt.name) and opt.bucket:
os.system(
"gsutil cp %s gs://%s/results/results%s.txt"
% (results_file, opt.bucket, opt.name)
)
# Log
tags = [
"train/box_loss",
"train/obj_loss",
"train/cls_loss", # train loss
"metrics/precision",
"metrics/recall",
"metrics/mAP_0.5",
"metrics/mAP_0.5:0.95",
"val/box_loss",
"val/obj_loss",
"val/cls_loss", # val loss
"x/lr0",
"x/lr1",
"x/lr2",
] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if tb_writer:
tb_writer.add_scalar(tag, x, epoch) # tensorboard
if wandb:
wandb.log({tag: x}) # W&B
# Update best mAP
fi = fitness(
np.array(results).reshape(1, -1)
) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, "r") as f: # create checkpoint
ckpt = {
"epoch": epoch,
"best_fitness": best_fitness,
"training_results": f.read(),
"model": ema.ema,
"optimizer": None if final_epoch else optimizer.state_dict(),
"wandb_id": wandb_run.id if wandb else None,
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
n = opt.name if opt.name.isnumeric() else ""
fresults, flast, fbest = (
save_dir / f"results{n}.txt",
wdir / f"last{n}.pt",
wdir / f"best{n}.pt",
)
for f1, f2 in zip(
[wdir / "last.pt", wdir / "best.pt", results_file], [flast, fbest, fresults]
):
if f1.exists():
os.rename(f1, f2) # rename
if str(f2).endswith(".pt"): # is *.pt
strip_optimizer(f2) # strip optimizer
os.system(
"gsutil cp %s gs://%s/weights" % (f2, opt.bucket)
) if opt.bucket else None # upload
# Finish
if plots:
plot_results(save_dir=save_dir) # save as results.png
if wandb:
files = [
"results.png",
"precision_recall_curve.png",
"confusion_matrix.png",
]
wandb.log(
{
"Results": [
wandb.Image(str(save_dir / f), caption=f)
for f in files
if (save_dir / f).exists()
]
}
)
logger.info(
"%g epochs completed in %.3f hours.\n"
% (epoch - start_epoch + 1, (time.time() - t0) / 3600)
)
else:
dist.destroy_process_group()
wandb.run.finish() if wandb and wandb.run else None
torch.cuda.empty_cache()
return results
def detect(save_img=False):
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# Initialize
device = torch_utils.select_device(opt.device)
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half = device.type != 'cpu' # half precision only supported on CUDA
# Load model
google_utils.attempt_download(weights)
model = torch.load(weights,
map_location=device)['model'].float() # load to FP32
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
# model.fuse()
model.to(device).eval()
imgsz = check_img_size(imgsz,
s=model.model[-1].stride.max()) # check img_size
if half:
model.half() # to FP16
# Second-stage classifier
classify = False
if classify:
modelc = torch_utils.load_classifier(name='resnet101',
n=2) # initialize
modelc.load_state_dict(
torch.load('weights/resnet101.pt',
map_location=device)['model']) # load weights
modelc.to(device).eval()
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = True
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz)
else:
save_img = True
dataset = LoadImages(source, img_size=imgsz)
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(len(names))]
# Run inference
t0 = time.time()
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
t1 = torch_utils.time_synchronized()
pred = model(img, augment=opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
classes=opt.classes,
agnostic=opt.agnostic_nms)
t2 = torch_utils.time_synchronized()
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
else:
p, s, im0 = path, '', im0s
save_path = str(Path(out) / Path(p).name)
txt_path = str(Path(out) / Path(p).stem) + (
'_%g' % dataset.frame if dataset.mode == 'video' else '')
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += '%g %ss, ' % (n, names[int(c)]) # add to string
# Write results
for *xyxy, conf, cls in det:
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 5 + '\n') %
(cls, *xywh)) # label format
if save_img or view_img: # Add bbox to image
label = '%s %.2f' % (names[int(cls)], conf)
plot_one_box(xyxy,
im0,
label=label,
color=colors[int(cls)],
line_thickness=3)
# Print time (inference + NMS)
print('%sDone. (%.3fs)' % (s, t2 - t1))
# Stream results
if view_img:
cv2.imshow(p, im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_img:
if dataset.mode == 'images':
cv2.imwrite(save_path, im0)
else:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release(
) # release previous video writer
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*opt.fourcc),
fps, (w, h))
vid_writer.write(im0)
if save_txt or save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
print('Done. (%.3fs)' % (time.time() - t0))
def train(hyp, opt, device, tb_writer=None, wandb=None):
logger.info(f'Hyperparameters {hyp}')
save_dir, epochs, batch_size, total_batch_size, weights, rank = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
# Directories
wdir = save_dir / 'weights'
wdir.mkdir(parents=True, exist_ok=True) # make dir
last = wdir / 'last.pt'
best = wdir / 'best.pt'
results_file = save_dir / 'results.txt'
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# Configure
plots = not opt.evolve # create plots
cuda = device.type != 'cpu'
init_seeds(2 + rank)
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict
with torch_distributed_zero_first(rank):
check_dataset(data_dict) # check
train_path = data_dict['train']
test_path = data_dict['val']
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(rank):
attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
if hyp.get('anchors'):
ckpt['model'].yaml['anchors'] = round(hyp['anchors']) # force autoanchor
model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device) # create
exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [] # exclude keys
state_dict = ckpt['model'].float().state_dict() # to FP32
state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect
model.load_state_dict(state_dict, strict=False) # load
logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report
else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for k, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
if opt.adam:
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum
else:
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
del pg0, pg1, pg2
# Scheduler https://arxiv.org/pdf/1812.01187.pdf
# https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf'] # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# plot_lr_scheduler(optimizer, scheduler, epochs)
# Logging
if wandb and wandb.run is None:
opt.hyp = hyp # add hyperparameters
wandb_run = wandb.init(config=opt, resume="allow",
project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
name=save_dir.stem,
id=ckpt.get('wandb_id') if 'ckpt' in locals() else None)
loggers = {'wandb': wandb} # loggers dict
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# Results
if ckpt.get('training_results') is not None:
with open(results_file, 'w') as file:
file.write(ckpt['training_results']) # write results.txt
# Epochs
start_epoch = ckpt['epoch'] + 1
if opt.resume:
assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
if epochs < start_epoch:
logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
(weights, ckpt['epoch'], epochs))
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, state_dict
# Image sizes
gs = int(max(model.stride)) # grid size (max stride)
imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples
# DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and rank != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()')
# EMA
ema = ModelEMA(model) if rank in [-1, 0] else None
# DDP mode
if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank)
# Trainloader
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
world_size=opt.world_size, workers=opt.workers,
image_weights=opt.image_weights)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
# Process 0
if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates
testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt, # testloader
hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True,
rank=-1, world_size=opt.world_size, workers=opt.workers, pad=0.5)[0]
if not opt.resume:
labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
if plots:
plot_labels(labels, save_dir, loggers)
if tb_writer:
tb_writer.add_histogram('classes', c, 0)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
# Model parameters
hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Start training
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
logger.info('Image sizes %g train, %g test\n'
'Using %g dataloader workers\nLogging results to %s\n'
'Starting training for %g epochs...' % (imgsz, imgsz_test, dataloader.num_workers, save_dir, epochs))
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if rank in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
# Broadcast if DDP
if rank != -1:
indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if rank != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(4, device=device) # mean losses
if rank != -1:
dataloader.sampler.set_epoch(epoch)
pbar = enumerate(dataloader)
logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'targets', 'img_size'))
if rank in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float() / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device), model) # loss scaled by batch_size
if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode
# Backward
scaler.scale(loss).backward()
# Optimize
if ni % accumulate == 0:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
# Print
if rank in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB)
s = ('%10s' * 2 + '%10.4g' * 6) % (
'%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
pbar.set_description(s)
# Plot
if plots and ni < 3:
f = save_dir / f'train_batch{ni}.jpg' # filename
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
# if tb_writer:
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
elif plots and ni == 3 and wandb:
wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg')]})
# end batch ------------------------------------------------------------------------------------------------
# end epoch ----------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard
scheduler.step()
# DDP process 0 or single-GPU
if rank in [-1, 0]:
# mAP
if ema:
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
final_epoch = epoch + 1 == epochs
if not opt.notest or final_epoch: # Calculate mAP
results, maps, times = test.test(opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
model=ema.ema,
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
plots=plots and final_epoch,
log_imgs=opt.log_imgs if wandb else 0)
# Write
with open(results_file, 'a') as f:
f.write(s + '%10.4g' * 7 % results + '\n') # P, R, [email protected], [email protected], val_loss(box, obj, cls)
if len(opt.name) and opt.bucket:
os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
# Log
tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
'x/lr0', 'x/lr1', 'x/lr2'] # params
for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
if tb_writer:
tb_writer.add_scalar(tag, x, epoch) # tensorboard
if wandb:
wandb.log({tag: x}) # W&B
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
# Save model
save = (not opt.nosave) or (final_epoch and not opt.evolve)
if save:
with open(results_file, 'r') as f: # create checkpoint
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'training_results': f.read(),
'model': ema.ema,
'optimizer': None if final_epoch else optimizer.state_dict(),
'wandb_id': wandb_run.id if wandb else None}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
# end epoch ----------------------------------------------------------------------------------------------------
# end training
if rank in [-1, 0]:
# Strip optimizers
final = best if best.exists() else last # final model
for f in [last, best]:
if f.exists():
strip_optimizer(f) # strip optimizers
if opt.bucket:
os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload
# Plots
if plots:
plot_results(save_dir=save_dir) # save as results.png
if wandb:
files = ['results.png', 'precision_recall_curve.png', 'confusion_matrix.png']
wandb.log({"Results": [wandb.Image(str(save_dir / f), caption=f) for f in files
if (save_dir / f).exists()]})
if opt.log_artifacts:
wandb.log_artifact(artifact_or_path=str(final), type='model', name=save_dir.stem)
# Test best.pt
logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
if opt.data.endswith('coco.yaml') and nc == 80: # if COCO
for conf, iou, save_json in ([0.25, 0.45, False], [0.001, 0.65, True]): # speed, mAP tests
results, _, _ = test.test(opt.data,
batch_size=total_batch_size,
imgsz=imgsz_test,
conf_thres=conf,
iou_thres=iou,
model=attempt_load(final, device).half(),
single_cls=opt.single_cls,
dataloader=testloader,
save_dir=save_dir,
save_json=save_json,
plots=False)
else:
dist.destroy_process_group()
wandb.run.finish() if wandb and wandb.run else None
torch.cuda.empty_cache()
return results
def train(
hyp, # path/to/hyp.yaml or hyp dictionary
opt,
device,
):
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
opt.resume, opt.noval, opt.nosave, opt.workers
# Directories
w = save_dir / 'weights' # weights dir
w.mkdir(parents=True, exist_ok=True) # make dir
last, best = w / 'last.pt', w / 'best.pt'
# Hyperparameters
if isinstance(hyp, str):
with open(hyp) as f:
hyp = yaml.safe_load(f) # load hyps dict
LOGGER.info(
colorstr('hyperparameters: ') + ', '.join(f'{k}={v}'
for k, v in hyp.items()))
# Save run settings
with open(save_dir / 'hyp.yaml', 'w') as f:
yaml.safe_dump(hyp, f, sort_keys=False)
with open(save_dir / 'opt.yaml', 'w') as f:
yaml.safe_dump(vars(opt), f, sort_keys=False)
# Config
plots = not evolve # create plots
cuda = device.type != 'cpu'
init_seeds(1 + RANK)
with open(data) as f:
data_dict = yaml.safe_load(f) # data dict
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
names = ['item'] if single_cls and len(
data_dict['names']) != 1 else data_dict['names'] # class names
assert len(
names
) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check
is_coco = data.endswith('coco.yaml') and nc == 80 # COCO dataset
# Loggers
if RANK in [-1, 0]:
loggers = Loggers(save_dir, weights, opt, hyp, data_dict,
LOGGER).start() # loggers dict
if loggers.wandb and resume:
weights, epochs, hyp, data_dict = opt.weights, opt.epochs, opt.hyp, loggers.wandb.data_dict
# Model
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(RANK):
weights = attempt_download(
weights) # download if not found locally
ckpt = torch.load(weights, map_location=device) # load checkpoint
model = Model(cfg or ckpt['model'].yaml,
ch=3,
nc=nc,
anchors=hyp.get('anchors')).to(device) # create
exclude = [
'anchor'
] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
csd = ckpt['model'].float().state_dict(
) # checkpoint state_dict as FP32
csd = intersect_dicts(csd, model.state_dict(),
exclude=exclude) # intersect
model.load_state_dict(csd, strict=False) # load
LOGGER.info(
f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}'
) # report
else:
model = Model(cfg, ch=3, nc=nc,
anchors=hyp.get('anchors')).to(device) # create
with torch_distributed_zero_first(RANK):
check_dataset(data_dict) # check
train_path, val_path = data_dict['train'], data_dict['val']
# Freeze
freeze = [] # parameter names to freeze (full or partial)
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
if any(x in k for x in freeze):
print(f'freezing {k}')
v.requires_grad = False
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size),
1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
g0, g1, g2 = [], [], [] # optimizer parameter groups
for v in model.modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
g2.append(v.bias)
if isinstance(v, nn.BatchNorm2d): # weight with decay
g0.append(v.weight)
elif hasattr(v, 'weight') and isinstance(
v.weight, nn.Parameter): # weight without decay
g1.append(v.weight)
if opt.adam:
optimizer = Adam(g0, lr=hyp['lr0'],
betas=(hyp['momentum'],
0.999)) # adjust beta1 to momentum
else:
optimizer = SGD(g0,
lr=hyp['lr0'],
momentum=hyp['momentum'],
nesterov=True)
optimizer.add_param_group({
'params': g1,
'weight_decay': hyp['weight_decay']
}) # add g1 with weight_decay
optimizer.add_param_group({'params': g2}) # add g2 (biases)
LOGGER.info(
f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "
f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias")
del g0, g1, g2
# Scheduler
if opt.linear_lr:
lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp[
'lrf'] # linear
else:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
scheduler = lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA
ema = ModelEMA(model) if RANK in [-1, 0] else None
# Resume
start_epoch, best_fitness = 0, 0.0
if pretrained:
# Optimizer
if ckpt['optimizer'] is not None:
optimizer.load_state_dict(ckpt['optimizer'])
best_fitness = ckpt['best_fitness']
# EMA
if ema and ckpt.get('ema'):
ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
ema.updates = ckpt['updates']
# Epochs
start_epoch = ckpt['epoch'] + 1
if resume:
assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.'
if epochs < start_epoch:
LOGGER.info(
f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs."
)
epochs += ckpt['epoch'] # finetune additional epochs
del ckpt, csd
# Image sizes
gs = max(int(model.stride.max()), 32) # grid size (max stride)
nl = model.model[
-1].nl # number of detection layers (used for scaling hyp['obj'])
imgsz = check_img_size(opt.imgsz, gs,
floor=gs * 2) # verify imgsz is gs-multiple
# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
logging.warning(
'DP not recommended, instead use torch.distributed.run for best DDP Multi-GPU results.\n'
'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.'
)
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
# Trainloader
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=opt.cache_images,
rect=opt.rect,
rank=RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '))
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(train_loader) # number of batches
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# Process 0
if RANK in [-1, 0]:
val_loader = create_dataloader(val_path,
imgsz,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=opt.cache_images and not noval,
rect=True,
rank=-1,
workers=workers,
pad=0.5,
prefix=colorstr('val: '))[0]
if not resume:
labels = np.concatenate(dataset.labels, 0)
# c = torch.tensor(labels[:, 0]) # classes
# cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency
# model._initialize_biases(cf.to(device))
if plots:
plot_labels(labels, names, save_dir, loggers)
# Anchors
if not opt.noautoanchor:
check_anchors(dataset,
model=model,
thr=hyp['anchor_t'],
imgsz=imgsz)
model.half().float() # pre-reduce anchor precision
# DDP mode
if cuda and RANK != -1:
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
# Model parameters
hyp['box'] *= 3. / nl # scale to layers
hyp['cls'] *= nc / 80. * 3. / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640)**2 * 3. / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.class_weights = labels_to_class_weights(
dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Start training
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * nb),
1000) # number of warmup iterations, max(3 epochs, 1k iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
last_opt_step = -1
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0
) # P, R, [email protected], [email protected], val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = amp.GradScaler(enabled=cuda)
compute_loss = ComputeLoss(model) # init loss class
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
f'Using {train_loader.num_workers} dataloader workers\n'
f'Logging results to {save_dir}\n'
f'Starting training for {epochs} epochs...')
for epoch in range(
start_epoch, epochs
): # epoch ------------------------------------------------------------------
model.train()
# Update image weights (optional)
if opt.image_weights:
# Generate indices
if RANK in [-1, 0]:
cw = model.class_weights.cpu().numpy() * (
1 - maps)**2 / nc # class weights
iw = labels_to_image_weights(dataset.labels,
nc=nc,
class_weights=cw) # image weights
dataset.indices = random.choices(
range(dataset.n), weights=iw,
k=dataset.n) # rand weighted idx
# Broadcast if DDP
if RANK != -1:
indices = (torch.tensor(dataset.indices)
if RANK == 0 else torch.zeros(dataset.n)).int()
dist.broadcast(indices, 0)
if RANK != 0:
dataset.indices = indices.cpu().numpy()
# Update mosaic border
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(3, device=device) # mean losses
if RANK != -1:
train_loader.sampler.set_epoch(epoch)
pbar = enumerate(train_loader)
LOGGER.info(
('\n' + '%10s' * 7) %
('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))
if RANK in [-1, 0]:
pbar = tqdm(pbar, total=nb) # progress bar
optimizer.zero_grad()
for i, (
imgs, targets, paths, _
) in pbar: # batch -------------------------------------------------------------
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float(
) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(
1,
np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [
hyp['warmup_bias_lr'] if j == 2 else 0.0,
x['initial_lr'] * lf(epoch)
])
if 'momentum' in x:
x['momentum'] = np.interp(
ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5,
imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]
] # new shape (stretched to gs-multiple)
imgs = nn.functional.interpolate(imgs,
size=ns,
mode='bilinear',
align_corners=False)
# Forward
with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
loss, loss_items = compute_loss(
pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize
if ni - last_opt_step >= accumulate:
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
last_opt_step = ni
# Log
if RANK in [-1, 0]:
mloss = (mloss * i + loss_items) / (i + 1
) # update mean losses
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%10s' * 2 + '%10.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss,
targets.shape[0], imgs.shape[-1]))
loggers.on_train_batch_end(ni, model, imgs, targets, paths,
plots)
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
if RANK in [-1, 0]:
# mAP
loggers.on_train_epoch_end(epoch)
ema.update_attr(model,
include=[
'yaml', 'nc', 'hyp', 'names', 'stride',
'class_weights'
])
final_epoch = epoch + 1 == epochs
if not noval or final_epoch: # Calculate mAP
results, maps, _ = val.run(data_dict,
batch_size=batch_size //
WORLD_SIZE * 2,
imgsz=imgsz,
model=ema.ema,
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco and final_epoch,
verbose=nc < 50 and final_epoch,
plots=plots and final_epoch,
loggers=loggers,
compute_loss=compute_loss)
# Update best mAP
fi = fitness(np.array(results).reshape(
1, -1)) # weighted combination of [P, R, [email protected], [email protected]]
if fi > best_fitness:
best_fitness = fi
loggers.on_train_val_end(mloss, results, lr, epoch, best_fitness,
fi)
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch':
epoch,
'best_fitness':
best_fitness,
'model':
deepcopy(de_parallel(model)).half(),
'ema':
deepcopy(ema.ema).half(),
'updates':
ema.updates,
'optimizer':
optimizer.state_dict(),
'wandb_id':
loggers.wandb.wandb_run.id if loggers.wandb else None
}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
del ckpt
loggers.on_model_save(last, epoch, final_epoch, best_fitness,
fi)
# end epoch ----------------------------------------------------------------------------------------------------
# end training -----------------------------------------------------------------------------------------------------
if RANK in [-1, 0]:
LOGGER.info(
f'{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.\n'
)
if not evolve:
if is_coco: # COCO dataset
for m in [last, best
] if best.exists() else [last]: # speed, mAP tests
results, _, _ = val.run(
data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
model=attempt_load(m, device).half(),
iou_thres=
0.7, # NMS IoU threshold for best pycocotools results
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=True,
plots=False)
# Strip optimizers
for f in last, best:
if f.exists():
strip_optimizer(f) # strip optimizers
loggers.on_train_end(last, best, plots)
torch.cuda.empty_cache()
return results
